From 25bd3cdc3355ce93587a3aaa05d98597bb8e6270 Mon Sep 17 00:00:00 2001
From: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
Date: Wed, 23 Sep 2026 16:48:16 +0800
Subject: [PATCH 1/2] feat(hir): simulate device collectives [M0]
---
docs/plans/engine-distributed-check.md | 159 ++++++
docs/rfcs/engine-scope-optimization.md | 519 ++++++++++++++++++
docs/spec/architecture.md | 6 +
docs/spec/evaluator.md | 71 ++-
docs/spec/hir.md | 68 +++
docs/spec/shard.md | 3 +
src/tilefoundry/evaluator/context.py | 2 +
src/tilefoundry/evaluator/distributed.py | 269 +++++++++
src/tilefoundry/evaluator/interpreter.py | 80 ++-
src/tilefoundry/ir/hir/sharding/collective.py | 132 +++++
.../ir/types/shard/shard_layout.py | 4 +-
tests/evaluator/test_distributed.py | 181 ++++++
tests/fixtures/distributed/__init__.py | 1 +
tests/fixtures/distributed/projection.py | 40 ++
14 files changed, 1507 insertions(+), 28 deletions(-)
create mode 100644 docs/plans/engine-distributed-check.md
create mode 100644 docs/rfcs/engine-scope-optimization.md
create mode 100644 src/tilefoundry/evaluator/distributed.py
create mode 100644 src/tilefoundry/ir/hir/sharding/collective.py
create mode 100644 tests/evaluator/test_distributed.py
create mode 100644 tests/fixtures/distributed/__init__.py
create mode 100644 tests/fixtures/distributed/projection.py
diff --git a/docs/plans/engine-distributed-check.md b/docs/plans/engine-distributed-check.md
new file mode 100644
index 00000000..9d552fcf
--- /dev/null
+++ b/docs/plans/engine-distributed-check.md
@@ -0,0 +1,159 @@
+---
+type: FEAT
+component: distributed-hir
+target_repo: tilefoundry
+---
+
+# [FEAT][distributed-hir] Execute and check explicit device collectives
+
+## Description
+
+Implement the first executable slice of the engine-scope RFC: rank-local
+evaluation of distributed projections, explicit collectives, and candidate HIR
+comparison through the existing `check` command. Distributed costing and ragged
+expert exchange remain subsequent RFC work.
+
+### Current state
+
+- `src/tilefoundry/evaluator/interpreter.py:122` evaluates a mesh body once on logical tensors.
+- `src/tilefoundry/ir/hir/sharding/reshard.py:277` evaluates a layout change without communication.
+- `src/tilefoundry/ir/hir/nn/matmul.py:204` derives partial output ownership for split contractions.
+- `src/tilefoundry/cli/check.py:306` compares runtime twins or saved outputs, with no second HIR selection.
+- `src/tilefoundry/ir/types/shard/shard_layout.py:182` maps factored layout positions to logical axes.
+
+### Decisions
+
+- D1 Participant groups -- reuse a concrete single-level device Mesh and one mesh-axis attribute per collective; remaining coordinates identify independent groups. This preserves existing placement ownership.
+- D2 Primitive slice -- implement AllReduce, AllGather, and ReduceScatter over equal partitions; reject unsupported distributions explicitly. Ragged exchange requires a separate payload contract.
+- D3 Ordering -- use the existing deterministic operand-order DAG walk, with function calls and uniform structured loops inlined conceptually. All ranks execute each collective occurrence together; downstream execution must preserve that group order.
+- D4 Simulation -- add an opt-in distributed evaluator context with rank-local values, reuse local primitive handlers, and reconstruct complete outputs only at the public boundary. No automatic completion of Partial outputs or communication through Reshard.
+- D5 Supported local computation -- initially support pointwise arithmetic, MatMul, Cast, Transpose, tuple projection, and reductions over unsplit axes; fail closed on other device-local operations rather than pretend global evaluation is distributed execution.
+- D6 HIR comparison -- add --reference SOURCE and --distributed to check. Bind activations positionally with matching logical shape/dtype and constants by the same resource paths. Reject incompatible signatures or conflicting reference modes.
+- D7 Evidence -- add a portable projection/state example and numerical public tests; no new eval command, strategy registry, or composed-kernel cost function. Existing analysis must refuse collectives until communication costs are implemented.
+- D8 Partition support -- use canonical or unfactored contiguous partitions and one mesh split per logical tensor axis; other layouts fail explicitly because their gathering may require additional redistribution.
+- D9 Bound dimensions -- accept leading layout factors across mixed factored/static and unfactored/dynamic axes. Preserve zero-sized logical axes during factor-to-axis mapping; otherwise an empty shard's ownership shifts to another dimension. Refines D8.
+
+## Milestones
+
+### Milestone M0: Collective semantics and rank simulation
+
+#### Depends
+- None
+
+#### Target State Design
+
+##### Delivered
+```diff
+# src/tilefoundry/ir/hir/sharding/collective.py
++class AllReduce(Op): ...
++class AllGather(Op): ...
++class ReduceScatter(Op): ...
+# src/tilefoundry/evaluator/interpreter.py
+-def evaluate(target, *inputs):
++def evaluate(target, *inputs, distributed=False):
+# src/tilefoundry/evaluator/distributed.py
++class DistributedValue(Value): ...
++def distribute(value, type_, bindings): ...
++def reconstruct(value): ...
++def evaluate_distributed_op(ctx, handler): ...
+# src/tilefoundry/ir/types/shard/shard_layout.py
++# Pseudocode: consume factors of a zero-sized logical axis through its zero factor.
+```
+
+##### Accepted by
+
+New numerical projection/state tests exercise real DSL parsing, type inference,
+local evaluation and communication. Incorrect subgroup membership, hidden
+reductions, or a dropped rank breaks observable results. Existing evaluator,
+sharding and parser suites remain regression coverage unchanged.
+
+Validation: the combined evaluator, IR operation, mesh/parser, CLI and analysis
+run passed 466 tests on CPU. One existing cross-device check was skipped and
+one existing CLI test selecting CUDA by default was excluded after confirming
+that this machine has no CUDA device. Numerical cases include dynamic and empty
+partitions, function/loop boundaries, and repeated explicit state updates.
+
+- [x] Partitioned projections agree with independently calculated torch results.
+- [x] Multi-axis groups preserve separate replicas and partition order.
+- [x] Partial outputs and communication-requiring Reshard fail explicitly.
+- [x] State outputs can be supplied to a following invocation with matching results.
+- [x] Local evaluation without distributed mode preserves existing behavior.
+
+
+- [ ] Touched tests MUST be reviewed for redundancy: remove ones superseded by the retained workflow, and do not add source-shape or hypothetical-refactor guards unless that form is a public contract.
+- [ ] A milestone that changes a public contract MUST list the owning `docs/spec/*.md` path in its `#### Related Files`; one that changes none lists no spec path.
+
+
+#### Related Files
+- `src/tilefoundry/ir/hir/sharding/collective.py`
+- `src/tilefoundry/evaluator/`
+- `src/tilefoundry/ir/types/shard/shard_layout.py`
+- `tests/evaluator/test_distributed.py`
+- `tests/fixtures/distributed/projection.py`
+- `docs/spec/architecture.md`
+- `docs/spec/hir.md`
+- `docs/spec/evaluator.md`
+- `docs/spec/shard.md`
+
+### Milestone M1: Explicit reference HIR checking
+
+#### Depends
+- M0
+
+#### Target State Design
+
+##### Delivered
+```diff
+# src/tilefoundry/cli/check.py
+ def add_arguments(parser):
++ parser.add_argument("--reference", metavar="SOURCE")
++ parser.add_argument("--distributed", action="store_true")
+ def check_concrete(request):
++ # Pseudocode: bind the same logical inputs/resource to both selected HIRs.
++ reference = evaluate_reference if request.reference is not None else existing_reference
+# docs/tutorial/distributed-check.md
++# Check a distributed projection against its reference
+# docs/tutorial/distributed-check.ipynb
++# Executable source, extraction command, comparison command and recorded output.
+# src/tilefoundry/cli/tutorial.py
++PAGES = (*existing_pages, "distributed-check")
+```
+
+##### Accepted by
+
+Extend the existing CLI behavioral suite with explicit HIR selection and the
+distributed fixture. This is the public agent workflow; a false numerical pass,
+independently drawn weights, or an ambiguous reference selection breaks it.
+
+Validation: both distributed projections pass explicit-reference CLI checks;
+an altered reference fails numerically and incompatible weight paths are
+refused. The tutorial notebook was executed through the existing renderer,
+including source extraction, and its Markdown records the passing output.
+All four analysis selectors were also exercised on a two-device deployment;
+each refuses the unregistered collective cost evaluator explicitly. Ruff,
+specification, comment, reference, language and path checks passed.
+
+- [x] CLI compares reference and distributed HIR using caller-provided predicates.
+- [x] JSON reports identify both selected programs and the evaluation mode.
+- [x] Mismatched bindings and conflicting reference flags fail before comparison.
+- [x] Existing expected-file and runtime-twin checks retain their behavior.
+
+
+- [ ] Touched tests MUST be reviewed for redundancy: remove ones superseded by the retained workflow, and do not add source-shape or hypothetical-refactor guards unless that form is a public contract.
+- [ ] A milestone that changes a public contract MUST list the owning `docs/spec/*.md` path in its `#### Related Files`; one that changes none lists no spec path.
+
+
+#### Related Files
+- `src/tilefoundry/cli/check.py`
+- `src/tilefoundry/cli/tutorial.py`
+- `tests/cli/test_cli_check.py`
+- `tests/fixtures/distributed/projection.py`
+- `docs/spec/cli.md`
+- `docs/tutorial/distributed-check.md`
+- `docs/tutorial/distributed-check.ipynb`
+- `docs/tutorial/index.md`
+
+## Final Gate
+
+
+
diff --git a/docs/rfcs/engine-scope-optimization.md b/docs/rfcs/engine-scope-optimization.md
new file mode 100644
index 00000000..fef726e0
--- /dev/null
+++ b/docs/rfcs/engine-scope-optimization.md
@@ -0,0 +1,519 @@
+# RFC: Engine-scope optimization with distributed HIR
+
+Status: Accepted for implementation. The first implementation slice covers
+device collectives and explicit reference-HIR correctness checking;
+distributed cost analysis and broader strategy coverage remain on the roadmap.
+
+## 1. Summary
+
+Extend TileFoundry's agent workflow with an optimizer stage that transforms a
+validated reference HIR into a distributed HIR and evaluates its resource and
+performance costs. The objective is to maximize throughput under a supplied SLO
+and hardware budget.
+
+The optimizer has two substages: transformation and cost evaluation. The agent
+authors transformations using specifications and examples; `check` validates
+their numerical behavior, and `analyze` derives costs from the resulting HIR.
+The agent uses those reports to revise or select a candidate before entering the
+existing local optimization and runtime implementation workflow.
+
+This design is independent of a particular model, accelerator, communication
+library, or serving framework. Model sources, weights, hardware facts, workload
+points, and SLO values are inputs to an optimization run.
+
+The architectural acceptance criterion is:
+
+> An agent can author a new legal composition of supported HIR operations,
+> check its behavior, and analyze its costs without registering a transformation
+> rule or writing a cost function for the composed kernel or strategy.
+
+## 2. Motivation and current state
+
+The [current workflow](../tutorial/index.md) first establishes a correct HIR
+reference and then improves its implementation. Engine-level decisions add
+another source of performance variation: which devices own weights, activations,
+and persistent state; which computation they perform; and how they communicate.
+These decisions also determine whether a workload fits in memory.
+
+TileFoundry already provides useful foundations:
+
+| Foundation | Current contract and gap |
+| --- | --- |
+| [HIR](../spec/hir.md) and [sharding](../spec/shard.md) | Functions, structured loops, meshes, layouts, and partial reductions describe computation and ownership. Distributed communication needs explicit semantics. |
+| [Evaluator](../spec/evaluator.md) | Executes logical values on torch tensors. It does not currently simulate communication among mesh participants. |
+| [`check`](../spec/cli.md#check) | Compares implementation outputs with reference results using caller-supplied predicates. Explicit reference-HIR versus candidate-HIR comparison needs a distributed execution path. |
+| [Analysis](../spec/analysis.md) | Derives work, traffic, memory, and nominal timing from authored HIR. Distributed links, per-device persistent state, and communication scheduling need coverage. |
+| [Target facts](../spec/target.md) | Describe capabilities and device resources. A deployment also needs concrete connectivity, device counts, and link resources. |
+| Separately supplied parallelism examples | Explain distributed strategies using semantic pseudocode. Their workload constants and candidate mappings are examples, not universal optimizer constraints. |
+
+Two existing boundaries require particular care. `Reshard` currently treats a
+same-storage layout change as a view; this is not sufficient to express moving
+data between devices. The existing `performance` analysis describes a nominal
+local execution model; extending engine analysis must preserve the meaning of
+its existing results.
+
+## 3. Goals and scope
+
+The initial scope includes tensor parallelism, expert parallelism, expert tensor
+parallelism, data parallelism, and context parallelism for prefill and decode.
+Their combinations are expressed through HIR ownership, computation, and
+communication. Separate prefill and decode placements may exchange explicitly
+represented persistent state.
+
+The stage delivers a selected distributed HIR, its deployment mapping and state
+contracts, and the evidence used to select it. The existing workflow then
+implements and optimizes that program locally.
+
+The following boundaries keep the work focused:
+
+- The public workflow uses `check` and `analyze`; this RFC adds no `eval` command
+ and no compiler-owned strategy-search command.
+- Strategy documents and examples guide the agent. They do not enumerate the
+ only permitted transformations.
+- Cost analysis is compositional over supported HIR. Primitive operation models
+ and hardware facts remain necessary; a handwritten cost model per composed
+ kernel or parallel strategy is not required.
+- Timing estimates support candidate ranking and approximate SLO filtering.
+ Full serving-runtime simulation, queueing, admission control, continuous
+ batching, and packet-level network simulation are outside the initial scope.
+- Profiling and detailed kernel calibration are optional later refinements.
+ They are not prerequisites for evaluating a candidate.
+- Pipeline parallelism is deferred because it introduces additional scheduling
+ and activation-lifetime decisions. It should later be expressed through an
+ explicit execution model rather than treated as an ordinary tensor split.
+
+## 4. Workflow and ownership
+
+```mermaid
+flowchart TD
+ R[Validated reference HIR] --> T[Agent authors candidate distributed HIR]
+ I[Workload, SLO, hardware facts, specs and examples] --> T
+ T --> C[check: reference versus candidate]
+ C -->|Failure diagnostics| T
+ C -->|Agreement on declared cases| A[analyze: resources and estimated cost]
+ A --> S[Agent compares feasible candidates]
+ S -->|Next hypothesis| T
+ S -->|Selected candidate| L[Existing local optimization and implementation]
+ L -->|Updated HIR and evidence| T
+```
+
+Transformation includes choosing placements, sharding, replication, algorithms,
+and communication arrangements. Correctness checking is the acceptance gate for
+a transformation. Cost evaluation measures the authored candidate and does not
+silently rewrite it or select a strategy.
+
+The agent owns candidate generation, workload sweeps, search budget, ranking,
+and stopping. TileFoundry owns operation semantics, legality checks, numerical
+evaluation, hardware facts, and reproducible analysis. Cheap legality or capacity
+checks may reject a candidate before expensive numerical evaluation; promotion
+still requires both correctness evidence and a resource/cost report.
+
+Search produces the best feasible candidate found within its declared budget,
+not a claim of global optimality. If no candidate meets the constraints, report
+that outcome with the relevant bottlenecks instead of changing the workload or
+silently relaxing the SLO.
+
+## 5. Inputs and outputs
+
+### 5.1 Optimization inputs
+
+The following are conceptual records, not proposed Python classes or CLI syntax:
+
+| Input | Required information |
+| --- | --- |
+| Reference program | HIR entry points, logical inputs and outputs, weight bindings, persistent-state transitions, numerical comparison policy, and source provenance. |
+| Hardware deployment | Device capabilities, counts, memory capacities, compute and memory rates, connectivity, communication capabilities, and provenance of supplied facts. |
+| Workload | Execution phase, total context, uncached token count, per-replica batch points, dtypes, initial state, and any data-dependent routing assumptions or bounds. |
+| Objective | Throughput definition, latency budgets, hardware budget, required workload coverage, and any workload mixture weights. |
+| Search configuration | Evaluation budget, candidate history, and permitted implementation capabilities. |
+
+Workload dimensions have distinct meanings. For a contiguous reusable prefix,
+`total_context = reusable_prefix + uncached_tokens`; the uncached count includes
+all tokens that must be recomputed after cache lookup. Cached context still
+affects attention reads and persistent storage. Recurrent state and other model
+state are described by the program rather than inferred from a universal KV
+formula.
+
+Lengths, batch limits, hit rates, and SLO thresholds belong to workload profiles,
+not compiler semantics or strategy documentation. Inconsistent profiles are
+rejected. The physical deployment is supplied separately from instruction-set
+capabilities: an architecture identifier alone does not establish available
+memory, device count, or connectivity.
+
+### 5.2 Optimization outputs
+
+A selected candidate is delivered with:
+
+- The authored distributed HIR and its relationship to the reference.
+- Device/group mapping and input, weight, output, and persistent-state layouts.
+- The supported workload domain and selected batch points.
+- Correctness reports with the cases, predicates, and bounds actually checked.
+- Analysis reports with capacity, traffic, communication, timing assumptions,
+ and the reasons for selection.
+- Source, workload, and hardware identities sufficient to reproduce the result.
+
+A deployment manifest may reference HIR definitions and supply physical rank
+mapping. It must not become a second, independently editable definition of
+tensor ownership or communication semantics.
+
+## 6. Distributed HIR
+
+### 6.1 One authored program
+
+Distributed HIR remains ordinary HIR extended with the necessary primitives.
+Global tensor shapes state logical meaning; layouts and meshes state ownership.
+Programs for individual ranks are downstream execution or lowering products.
+Analysis may construct an internal event graph, but it remains a derived view of
+the authored HIR rather than another source language.
+
+Device-level distribution and finer execution levels remain distinguishable.
+The same physical devices may have different logical group views for different
+subgraphs. There is no universal requirement that parallel degrees multiply
+together or that an expert group equal a particular attention group.
+
+### 6.2 Communication semantics
+
+The initial primitive coverage needs broadcast, gather, scatter, all-gather,
+reduction, all-reduce, reduce-scatter, and fixed or variable-size all-to-all.
+Some may be expressible as compositions; the owner specs will settle the minimal
+primitive set. These names identify required semantics, not existing DSL calls.
+
+Each communication operation defines:
+
+- Participants and their rank order or coordinate mapping.
+- Payload shapes, dtypes, source and destination ownership, and reduction kind
+ where applicable.
+- Value dependencies and ordering relative to other communication in the same
+ group, including participation by ranks with no payload.
+- Counts, offsets, padding, and capacity bounds for variable-size exchanges.
+- Reference evaluation, access/traffic relations, legality, and requirements on
+ a downstream implementation.
+
+Cross-device movement is explicit. A layout annotation alone cannot cause an
+unreported transfer or convert a partial value into a complete value. Existing
+`Reshard` behavior remains unchanged initially; a future redistribution shorthand
+would need to elaborate into explicit communication with the same checks and
+cost accounting.
+
+Collectives must have a consistent ordering across participating ranks. The
+semantic design must specify how that ordering survives HIR traversal and later
+rewrites. This can use explicit dependencies or an ordering representation; the
+choice must be settled before collective implementation.
+
+### 6.3 Persistent and data-dependent state
+
+State remains explicit in function inputs and outputs. Examples include cache
+append, recurrent-state replacement, and bounded convolution windows. Prefix
+reuse supplies the state associated with the reused prefix boundary. If
+prefill and decode use different layouts, their boundary includes a state-layout
+conversion whose semantics, temporary storage, and traffic are visible.
+
+Expert dispatch preserves token identity, routing weights, and the information
+needed to restore output order. Numerical evaluation uses actual counts.
+Analysis uses declared counts, ranges, or workload assumptions and identifies
+which results depend on them. Balanced routing may be a timing assumption, but
+it cannot silently stand in for a capacity bound.
+
+Context-parallel examples must express the algorithm's actual dependency. For
+softmax attention this includes normalization-aware combination of partial
+results; for recurrent algorithms it includes propagation or composition of
+state across token partitions. These are executable HIR examples assembled from
+supported operations, not strategy names with hidden implementations.
+
+## 7. Correctness through `check`
+
+Extend `check` to accept an explicit reference HIR and candidate HIR while
+preserving existing runtime-twin and expected-output workflows. Both sides use
+the same logical activations, weights, dimensions, and initial state. Reference
+selection and input binding are part of the report. Exact argument spelling is
+left to the CLI contract work.
+
+The internal evaluator gains a mode that simulates logical ranks and their local
+tensors on one physical device. It executes local computation and collective
+semantics and reconstructs the declared observable outputs. It must not replace
+the candidate with the original logical computation or repair a missing
+collective implicitly.
+
+Comparison includes persistent-state outputs and repeated invocations that
+consume them. Different physical layouts are compared through their declared
+logical correspondence. A different state representation requires an explicit,
+checkable reconstruction contract. Evaluator-only gathering for comparison is
+not charged as deployed communication unless it is also part of the candidate.
+
+Numerical predicates and tolerances remain caller-supplied. Integral routing and
+index results use appropriate exact checks. A passing report states agreement
+on the supplied cases; it is not a proof for every possible input. Shape,
+placement, and collective legality checks complement numerical evaluation.
+
+Correctness cases may use smaller concrete sizes to exercise partition
+boundaries, empty routes, padding, and state transitions. Cost analysis still
+uses the declared production workload sizes. Reports distinguish these domains;
+an unsupported or unexecuted case never becomes a correctness pass.
+
+Semantic simulation does not validate a communication library or establish that
+physical execution is deadlock-free. That remains part of validating the runtime
+implementation in the downstream workflow.
+
+## 8. Compositional cost evaluation through `analyze`
+
+### 8.1 Derive costs from the candidate
+
+Extend the existing analysis machinery over the same typed HIR. Primitive
+operation semantics and access relations determine work and data movement;
+layouts determine local ownership and replication; target facts supply rates
+and capacities. Loops and function composition aggregate those quantities.
+
+A new composition of supported primitives requires no kernel-specific cost
+registration. A genuinely new primitive requires its semantic, access, and
+analysis contracts. Unsupported operations or missing required facts produce
+diagnostics or an explicitly incomplete report, never an implicit zero cost.
+
+Representative algorithm HIRs expose the implementation choices relevant to
+strategy ranking: tiled attention, online normalization, quantization and
+scales, cache writes, sparse reads, routing, and communication. Their names do
+not select opaque cost formulas. Tiling, fusion, or replication changes costs
+because the authored storage and accesses change.
+
+Inlining or summarization must avoid charging both a function and its expanded
+body. Any execution boundary relevant to costing must be explicit or derived
+from documented semantics; an arbitrary helper-function boundary does not imply
+a device launch. Structured loops should remain analyzable without materializing
+one graph node per token or iteration.
+
+### 8.2 Memory feasibility
+
+For each device and declared execution schedule, estimate the peak simultaneous
+allocation of resident weights, persistent state, live intermediates, and
+communication/workspace buffers, then include the declared runtime reserve.
+Persistent state remains live across the invocations that require it.
+
+Capacity and traffic are separate quantities. All resident expert weights count
+toward capacity even when an invocation reads only a subset. Retaining a cache
+does not mean every operation reads all of it. Shared storage, aliasing, in-place
+updates, and reused prefixes reduce capacity only when their ownership and
+lifetimes justify that reduction. The evaluator's temporary torch allocations
+are not the candidate's physical memory model.
+
+Variable-size operations need a stated allocation bound or capacity policy.
+Unknown bounds leave capacity unresolved. Fragmentation, reserved memory, and
+optional state snapshots are explicit deployment/workload inputs rather than
+hidden constants. State transfer includes temporary source and destination
+storage where both are live.
+
+The report distinguishes a conservative capacity result from an estimate that
+depends on an assumed allocation schedule. An optimistic footprint alone cannot
+establish feasibility.
+
+### 8.3 Communication and timing
+
+Communication costs derive from explicit payloads, participants, and topology.
+A generic latency-plus-transfer model is sufficient initially. Link bandwidth,
+startup latency, routes, and shared resources come from attributed target or
+deployment facts. Communication-buffer reads and writes are accounted for
+consistently with network transfers, avoiding both omissions and duplicate
+charges for the same access.
+
+Compose primitive costs over data and communication dependencies with a simple,
+documented resource model. The initial schedule can serialize work sharing a
+resource and permit overlap where dependencies and resource independence allow
+it. Refining contention or overlap later should preserve the reported model
+identity. Summing all devices' work is not a substitute for estimating the
+critical path.
+
+The existing local `performance` result keeps its current meaning. Distributed
+timing is exposed through clearly identified records within `analyze`, with the
+schema and selector details settled in the analysis contract. Optional
+ideal-overlap results are labeled bounds and remain separate from the estimate
+for the authored execution structure.
+
+The report states what latency includes. Prefill execution time can be compared
+with an allocated TTFT budget, but it does not implicitly include request
+queueing or unmodeled service overhead. Decode step time, average time per
+emitted token, and inter-output latency remain distinct when one invocation can
+emit multiple tokens. Any conversion uses explicit output/acceptance assumptions.
+
+### 8.4 Report requirements
+
+Reports should expose enough evidence for an agent to form its next hypothesis:
+
+- Work and traffic by operation, execution scope, and device.
+- Resident state and peak allocations, including the limiting device and live
+ buffers at a capacity failure.
+- Communication bytes and estimated time by group and topology resource.
+- Estimated critical path and compute, memory, or communication bottlenecks.
+- Workload bindings, hardware provenance, modeling assumptions, and unsupported
+ or unresolved quantities.
+
+Human-readable and machine-readable reports come from the same analysis
+results. Findings retain source/operation provenance so the agent can identify
+the HIR responsible for a bottleneck.
+
+## 9. Search, examples, and handoff
+
+The agent uses specification and example documents to propose transformations.
+Examples describe preconditions, logical behavior, state handling, and expected
+resource tradeoffs, and include runnable HIR. They remain starting points for
+composition rather than a whitelist of legal strategies.
+
+Candidate comparison maximizes the workload's throughput measure subject to
+correctness, supported implementation capabilities, capacity, and approximate
+latency constraints. Batch sweeps report the feasible region; a candidate that
+fails at a larger batch can remain useful at smaller batches. Pruning larger
+points after a failure requires a justified monotonicity assumption.
+
+Keep a small Pareto frontier when throughput and latency trade off. For multiple
+workload classes, either use an explicitly supplied mixture/objective or report
+separate frontiers. An average must not hide a required workload's SLO failure.
+The report identifies every comparison's device budget and work completed, so
+replication or speculative work cannot inflate useful throughput.
+
+The initial policy selects a strategy per declared execution phase or pool.
+If a candidate switches layouts during a stateful request, the transition must
+appear in its program and cost. Reference semantics, numerical policy, workload,
+hardware facts, and analysis assumptions remain fixed when comparing candidates.
+
+The selected HIR enters the existing local optimization workflow with its
+communication and state contracts intact. Later changes to layout, buffering,
+or computation require renewed checking and analysis. Measured results can
+motivate another agent iteration without making measurement a prerequisite for
+the engine-stage search.
+
+## 10. What to do next
+
+The following phases are an implementation roadmap. Exact APIs and source diffs
+belong in subsequent implementation plans. Each phase should update its owning
+specifications together with the corresponding implementation.
+
+### Phase 1: Settle distributed semantic contracts
+
+Deliver the participant/group representation, global/local shape rules, partial
+reduction transitions, collective ordering, variable-size payload representation,
+and persistent-state boundary. Define portable workload and deployment inputs.
+Resolve the open interface questions in Section 12 before implementing affected
+surfaces.
+
+Acceptance: small model-independent HIR examples can state a projection split,
+a context/state split, and token dispatch/combine without opaque strategy nodes.
+Their communication, state correspondence, and required capabilities are explicit.
+
+Related files: [architecture](../spec/architecture.md), [HIR](../spec/hir.md),
+[sharding](../spec/shard.md), [types](../spec/types.md),
+[parser](../spec/parser.md), and [target](../spec/target.md).
+
+### Phase 2: Enable distributed correctness checking
+
+Implement collective primitives and their type/layout checks, the rank-simulating
+evaluator, and explicit reference-versus-candidate selection through `check`.
+Preserve existing command behavior and numerical comparison policy.
+
+Acceptance: reference and candidate agree for representative partitioned
+computation and repeated state transitions; deliberately missing reductions,
+incorrect routing, or incompatible groups are detected. Cover externally
+reachable empty-route and uneven-partition cases where the contract supports
+them. Reuse existing evaluator and command workflows where possible.
+
+Related files: `src/tilefoundry/ir/hir/`, `src/tilefoundry/evaluator/`,
+`src/tilefoundry/cli/check.py`, [HIR](../spec/hir.md),
+[evaluator](../spec/evaluator.md), and [CLI](../spec/cli.md).
+
+### Phase 3: Add distributed resource and cost analysis
+
+Extend primitive access/cost analysis, deployment facts, per-device memory
+ownership, communication accounting, and dependency-based timing. Add report
+records and diagnostics without changing the meaning of existing local results.
+
+Acceptance: analytically tractable examples establish correct byte ownership,
+resident-versus-read distinctions, communication payloads, and limiting-device
+capacity. Changing an input deployment's capacity or link rate affects the
+appropriate result. Unknown facts are reported. New compositions of supported
+operations work without a composed-kernel cost registration.
+
+Related files: `src/tilefoundry/analysis/`, `src/tilefoundry/visitor_registry/`,
+`src/tilefoundry/target/`, `src/tilefoundry/inspection/`,
+`src/tilefoundry/cli/analyze.py`, [analysis](../spec/analysis.md),
+[target](../spec/target.md), [inspection](../spec/inspection.md), and
+[CLI](../spec/cli.md).
+
+### Phase 4: Publish specifications and executable examples for agents
+
+Provide portable HIR examples for tensor and expert parallelism, prefill/decode
+context partitioning, recurrent-state handling, and their useful combinations.
+Include representative attention, routing, quantization, and cache-update
+algorithms as analyzable HIR. Expose the material through the existing installed
+documentation/example discovery surfaces.
+
+Acceptance: examples use public supported operations, pass `check`, and produce
+cost reports. An agent can alter a composition and receive useful feedback from
+the same commands. Deployment and workload values are supplied externally.
+
+Related files: `parallelism/`, `docs/tutorial/`, `tests/fixtures/`,
+`src/tilefoundry/cli/`, [CLI](../spec/cli.md), and package-data configuration.
+
+### Phase 5: Demonstrate the optimizer workflow and downstream handoff
+
+Run an agent-guided search with separately supplied reference, workload, and
+deployment profiles. Retain candidate HIRs, correctness results, cost reports,
+and the selected frontier. Demonstrate the selected program entering the
+existing implementation workflow with its state and communication contracts.
+
+Acceptance: the search reproduces its selection under fixed inputs, reports
+infeasible cases, and delivers a reviewable HIR. At least one candidate combines
+supported primitives in a way that needs neither a new transformation rule nor
+a new composed-kernel cost function. A small executable realization validates
+the handoff; broad kernel tuning and serving benchmarks remain downstream work.
+
+Related files: `docs/tutorial/`, portable workflow examples, and the existing
+integration/runtime workflows used for the realization.
+
+## 11. Validation strategy
+
+Use focused, reusable behavioral coverage for new semantics. The key evidence
+is numerical agreement across distribution and state transitions, correct
+resource accounting on small independently calculable programs, and preservation
+of existing command behavior. Tests should exercise plausible failures at public
+boundaries rather than assert source shape or a particular internal design.
+
+Large workload points primarily exercise analysis scalability and resource
+feasibility; they need not all materialize full tensors in the reference
+interpreter. Report this distinction explicitly. Runtime validation verifies the
+implementation of a selected HIR and its collectives; analytical predictions
+remain predictions until measured.
+
+The RFC is complete as a design proposal when the workflow boundaries, semantic
+requirements, ownership, roadmap, and acceptance evidence are reviewable.
+Implementation completion requires the phase-specific evidence above.
+
+## 12. Open design questions
+
+These are generic implementation decisions, independent of the first model or
+hardware deployment:
+
+1. What is the smallest collective primitive set, and how are ordering and
+ participant groups represented without compromising existing HIR value
+ semantics?
+2. Which representation best expresses variable-size token exchanges and their
+ capacity bounds within the current type and symbolic-dimension system?
+3. How should `check` name the second HIR and bind logical state or explicit
+ reconstruction adapters while preserving its existing input conventions?
+4. How should deployment connectivity extend target facts, and which analysis
+ record/schema changes expose distributed timing and capacity assumptions?
+5. How should workload profiles and candidate reports be serialized and exposed
+ through the existing command surfaces?
+
+The initial answers should favor the smallest coherent extensions to existing
+owners. They do not reopen the decisions to use agent-authored transformations,
+`check` for correctness, and compositional `analyze` for cost evaluation.
+
+## 13. Related work
+
+[RoofLang](https://arxiv.org/pdf/2609.12551v1) combines graph transformations,
+placement, and analytical simulation to guide architecture search. It provides
+useful examples of communication simplification, memory/traffic separation, and
+cost-based exploration. Its [implementation](https://github.com/yzygitzh/rooflang)
+uses explicit transformation helpers and kernel-level analytical descriptions.
+
+This proposal retains TileFoundry's executable HIR as the common foundation:
+the agent authors transformations from specs and examples, the evaluator checks
+their behavior, and analysis derives costs from their supported primitive
+composition. The initial cost model aims to expose architectural tradeoffs with
+declared assumptions rather than reproduce a particular serving stack.
diff --git a/docs/spec/architecture.md b/docs/spec/architecture.md
index 10c8cea3..5935df28 100644
--- a/docs/spec/architecture.md
+++ b/docs/spec/architecture.md
@@ -119,6 +119,12 @@ objects — DOT, Python DSL pretty-printer, dump integration, and the
interactive viewer — and never introduces new semantic ownership.
Concrete presentation contracts live in [inspection](./inspection.md).
+Distributed programs use the same HIR, with explicit device collectives owned
+by [HIR](./hir.md#3-device-collectives) and ownership carried by shard layouts.
+The [evaluator](./evaluator.md#7-distributed-evaluation) can simulate local tensors
+and communication on one torch device. The [CLI](./cli.md#check) compares that
+program with an explicitly selected reference HIR.
+
## 5. Analysis & optimization
This stage layers two concerns on top of the same IR:
diff --git a/docs/spec/evaluator.md b/docs/spec/evaluator.md
index 8a8cd0bd..46b32bd5 100644
--- a/docs/spec/evaluator.md
+++ b/docs/spec/evaluator.md
@@ -14,6 +14,7 @@ flowchart TB
Value["Value"]
TensorValue["TensorValue
(data, type)"]
TupleValue["TupleValue
(elements)"]
+ DistributedValue["DistributedValue
(shards, type, mesh)"]
evaluate --> Evaluator
Evaluator -. "Call(target=Op)" .-> registry
@@ -22,12 +23,14 @@ flowchart TB
Evaluator -. produces .-> Value
Value --> TensorValue
Value --> TupleValue
+ Value --> DistributedValue
```
```python
def evaluate(
target: "Function | LoadedModule",
*inputs: "torch.Tensor",
+ distributed: bool = False,
) -> "torch.Tensor | tuple[torch.Tensor, ...]":
...
```
@@ -47,7 +50,8 @@ default.
## 1. `Value`
The values that flow through evaluation form a small hierarchy: a
-single-output node produces a `TensorValue`; a multi-output node (a
+single-output node produces a `TensorValue`, or a `DistributedValue` under
+[distributed evaluation](#7-distributed-evaluation); a multi-output node (a
`Tuple`, a `TupleType` `Call`, or a multi-carry `LoopRegion`)
produces a `TupleValue`.
@@ -56,7 +60,8 @@ class Value:
"""Provide the base of every evaluated value."""
```
-- constraints: none — abstract base; concrete values are `TensorValue` / `TupleValue`
+- constraints: none — abstract base; concrete values are `TensorValue`,
+ `TupleValue`, and `DistributedValue`.
### `TensorValue`
@@ -154,6 +159,8 @@ class EvaluateContext:
loaded_module: attribute; Runtime module reading, when one is active.
device: attribute; Where the inputs are, or None to leave it to torch.
dim_bindings: attribute; concrete values for symbolic ShapeDims.
+ distributed: attribute; Whether to simulate device participants.
+ mesh: attribute; Device mesh active in the current evaluation scope.
"""
op: Any = None
@@ -162,6 +169,8 @@ class EvaluateContext:
loaded_module: Any | None = None
device: str | None = None
dim_bindings: Mapping[str, int] = field(default_factory=dict)
+ distributed: bool = False
+ mesh: "Mesh | None" = None
def for_op(self, op: Any, args: tuple[Any, ...], result_type: Any) -> EvaluateContext: ...
@@ -238,8 +247,8 @@ separate `eval_grid` function.
## 6. Layout domain
-Evaluation models a **single mesh participant** and operates on logical
-values:
+Ordinary evaluation (`distributed=False`) models a **single mesh participant**
+and operates on logical values:
- An axis-bearing op (`Reduce`, `rms_norm`, …) addresses its `axis` /
`axes` in the operand's **logical** `TensorType.shape`, regardless of
@@ -264,3 +273,57 @@ values:
raise `EvalError` saying that the evaluator models one mesh participant and
linking to this section, until the mesh evaluator models that path. An
unmodelled path MUST identify itself rather than leaking a torch exception.
+
+## 7. Distributed evaluation
+
+`evaluate(..., distributed=True)` executes device-local tensors and explicit
+[collectives](./hir.md#3-device-collectives) on the device holding the supplied
+torch inputs. It returns reconstructed logical outputs using the same public
+return structure as ordinary evaluation.
+
+```python
+class DistributedValue(Value):
+ """Hold each participant's tensor with the global ownership type."""
+
+ shards: tuple[torch.Tensor, ...]
+ type: TensorType
+ mesh: Mesh
+```
+
+- constraints:
+ - Meshes MUST be concrete, contiguous, single-level `gpu` meshes. Shards use
+ lexicographic coordinate order, with the last mesh axis varying fastest.
+ Nested different meshes, sliced meshes and mixed device/thread meshes MUST
+ be refused. The topology name identifies device ownership; simulation does
+ not require physical GPUs.
+ - Distributed ownership MUST use `Broadcast`, `Split` or `Partial`. A logical
+ axis MAY be split by one mesh axis. Equal contiguous partitions MUST be
+ represented by splits at the leading nonunit layout factor of each logical axis;
+ other partitions MUST be refused. Canonical factorization and dynamically
+ bound unfactored axes satisfy this rule. Bound symbolic dimensions MUST satisfy
+ runtime divisibility and shape checks.
+ - Logical inputs with declared `Split` ownership are partitioned at the entry
+ boundary; `Broadcast` inputs are replicated. These are input placement,
+ not communication operations. Logical inputs MUST NOT declare `Partial`:
+ a complete tensor cannot specify each rank's independent contribution.
+ - Unplaced inputs/constants supply the same complete value to each participant.
+ Inside a device region, `Reshard` MAY retain local ownership or select a
+ subset already present on each participant. It MUST refuse a destination
+ requiring another participant's data, creating partial contributions or
+ completing a partial reduction.
+ - `Binary`, `Unary`, `MatMul`, `Cast`, `Transpose`, and `Reduce` reuse their
+ registered evaluators with local shapes. A `Reduce` over an axis carrying
+ `Split` ownership MUST be refused. Tuple construction/projection, calls and
+ uniform structured loops preserve distributed values. Other device-local
+ operations MUST report unsupported semantics rather than execute their
+ global logical operation.
+ - Every returned `Split` tensor is assembled by its declared ownership.
+ `Broadcast` replicas MUST agree exactly, treating matching NaNs as equal.
+ Returning `Partial` MUST fail and name the missing explicit reduction.
+ Reconstruction is an evaluator observation boundary; it inserts no HIR
+ collective and establishes no runtime transfer cost.
+ - Persistent state uses ordinary returned tensors and subsequent inputs.
+ There is no hidden state store. A caller MAY feed reconstructed state into
+ the next invocation, where the declared input placement applies again.
+ - Without `distributed=True`, existing logical evaluation semantics remain
+ unchanged and device collectives MUST raise `EvalError`.
diff --git a/docs/spec/hir.md b/docs/spec/hir.md
index 462dc927..55815233 100644
--- a/docs/spec/hir.md
+++ b/docs/spec/hir.md
@@ -1694,3 +1694,71 @@ def is_concrete(fn: Function) -> bool:
and called functions. `is_concrete` additionally checks the return type and
is false exactly when a reachable required extent still contains a `DimVar`;
runtime values without a `DimVar` are rejected later by their consumer.
+
+## 3. Device collectives
+
+Device communication is explicit value-producing HIR. These operations live in
+`ir/hir/sharding/collective.py` and use the input's `ShardLayout.mesh`. A selected
+mesh axis varies within each group; all other coordinates remain fixed. Thus a
+two-axis mesh can express independent reduction groups without a strategy node.
+
+```python
+class AllReduce(Op):
+ """Reduce partial values and broadcast the complete result within each group."""
+
+ x: Tensor
+ mesh_axis: int
+
+class AllGather(Op):
+ """Concatenate equal partitions in increasing mesh-coordinate order."""
+
+ x: Tensor
+ mesh_axis: int
+
+class ReduceScatter(Op):
+ """Reduce partial values and partition the result equally within each group."""
+
+ x: Tensor
+ mesh_axis: int
+ tensor_axis: int
+```
+
+- constraints:
+ - The input MUST carry a `ShardLayout` whose mesh is bound by the current
+ `MeshRegion`. The mesh MUST have one concrete `gpu` topology and positive
+ concrete extents with contiguous C-order positions within that topology.
+ `mesh_axis` MUST be a nonnegative integer indexing this mesh.
+ - Each mesh axis MUST have a `Broadcast`, `Split` or `Partial` attribute. Each
+ logical tensor axis MAY be split by at most one mesh axis. Partitions MUST
+ be equal and contiguous, splitting at the leading nonunit layout factor of the
+ logical axis; nondivisible concrete shapes MUST fail.
+ - `AllReduce` and `ReduceScatter` MUST consume `Partial("sum")`,
+ `Partial("max")` or `Partial("min")` on `mesh_axis`; that attribute supplies
+ the reduction operation. `AllReduce` changes it to `Broadcast`.
+ - `AllGather` MUST consume `Split` on `mesh_axis` and changes it to
+ `Broadcast`. It preserves logical shape rather than introducing a rank
+ dimension or multiplying the global extent.
+ - `ReduceScatter.tensor_axis` MUST index the logical tensor shape and MUST
+ NOT already be split by another mesh axis. The result replaces the selected
+ `Partial` with a `Split` of that logical axis.
+ - All three preserve global logical shape, dtype, storage and ownership on
+ other mesh axes. Output layouts use canonical factorization; attributes
+ referring to unaffected logical axes are remapped accordingly.
+ - Every participant in a group participates in each collective occurrence,
+ including when a payload contains zero elements. Calls have value semantics:
+ they produce results without mutating their inputs.
+ - Collective occurrence order is the deterministic depth-first traversal of
+ operands in their declared order, with shared DAG nodes evaluated once per
+ execution scope. Function bodies and uniform `LoopRegion` iterations expand
+ in that order. All participants execute that same order. A runtime lowering
+ MUST preserve it within each group, and MUST NOT introduce rank-dependent
+ participation or reorder collectives independently on different ranks.
+ - Numerical evaluation requires [distributed evaluation](./evaluator.md#7-distributed-evaluation).
+ Ordinary logical evaluation MUST refuse these operations. Numerical simulation
+ checks value semantics; it does not verify a physical communication backend.
+ - Logical tensor-coordinate access is identity: every output index consumes
+ that same index in the input. Rank participation and ownership changes are
+ defined by the collective contract above. These operations have no registered
+ cost evaluator. An analysis needing that evaluator MUST report them as
+ unsupported, never as zero-cost layout views. The existing `Reshard` cost
+ classification is unchanged.
diff --git a/docs/spec/shard.md b/docs/spec/shard.md
index 8ab3d31b..15b6478b 100644
--- a/docs/spec/shard.md
+++ b/docs/spec/shard.md
@@ -602,6 +602,9 @@ Let `sl: ShardLayout`, `T: TensorType`, and `G = sl.layout.shape`.
- `local_shape(sl)[k] = G[k]` otherwise.
- The canonical regroup rule ([§8](#8-layout-propagation)) defines how `T.shape` aligns with
`G`.
+- Mapping layout factors to a zero-length logical axis MUST consume factors
+ through the first zero factor. That zero axis MUST retain its factors rather
+ than assigning their `Split` ownership to the following logical axis.
#### 7.1.2 `layout.strides`
diff --git a/src/tilefoundry/evaluator/context.py b/src/tilefoundry/evaluator/context.py
index 8e83ab08..2434b370 100644
--- a/src/tilefoundry/evaluator/context.py
+++ b/src/tilefoundry/evaluator/context.py
@@ -16,6 +16,8 @@ class EvaluateContext:
loaded_module: Any | None = None
device: str | None = None
dim_bindings: Mapping[str, int] = field(default_factory=dict)
+ distributed: bool = False
+ mesh: "Mesh | None" = None
def for_op(self, op: Any, args: tuple[Any, ...], result_type: Any) -> EvaluateContext:
"""Add one Call's evaluated operands while preserving runtime state."""
diff --git a/src/tilefoundry/evaluator/distributed.py b/src/tilefoundry/evaluator/distributed.py
new file mode 100644
index 00000000..349514b2
--- /dev/null
+++ b/src/tilefoundry/evaluator/distributed.py
@@ -0,0 +1,269 @@
+"""Simulate device-local tensors and explicit communication on one torch device."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from itertools import product
+
+import torch
+
+from tilefoundry.evaluator.dim import resolve_dim
+from tilefoundry.evaluator.value import EvalError, TensorValue, TupleValue, Value
+from tilefoundry.ir.hir.math.binary import Binary
+from tilefoundry.ir.hir.math.unary import Unary
+from tilefoundry.ir.hir.nn.matmul import MatMul
+from tilefoundry.ir.hir.sharding.collective import (
+ AllGather,
+ AllReduce,
+ ReduceScatter,
+ device_mesh_shape,
+)
+from tilefoundry.ir.hir.sharding.reshard import Reshard
+from tilefoundry.ir.hir.tensor.cast import Cast
+from tilefoundry.ir.hir.tensor.reduce import Reduce
+from tilefoundry.ir.hir.tensor.transpose import Transpose
+from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem
+from tilefoundry.ir.types import TensorType
+from tilefoundry.ir.types.shard import Layout, Mesh
+from tilefoundry.ir.types.shard.shard_layout import (
+ Broadcast,
+ Partial,
+ ShardLayout,
+ Split,
+ layout_axis_to_tensor_axis,
+)
+from tilefoundry.ir.types.substitute import substitute_dims
+
+
+@dataclass(frozen=True)
+class DistributedValue(Value):
+ """Hold concrete local tensors in lexicographic mesh-coordinate order."""
+
+ shards: tuple[torch.Tensor, ...]
+ type: TensorType
+ mesh: Mesh
+
+
+def coordinates(mesh):
+ """Enumerate independent participants, with the last mesh axis fastest."""
+ return tuple(product(*(range(size) for size in device_mesh_shape(mesh))))
+
+
+def _shape(type_, bindings):
+ return tuple(resolve_dim(size, bindings) for size in type_.shape)
+
+
+def _attrs(type_, mesh):
+ layout = type_.layout
+ if layout is None or isinstance(layout, Layout):
+ return (Broadcast(),) * len(mesh.layout.shape)
+ if not isinstance(layout, ShardLayout) or layout.mesh != mesh:
+ raise EvalError("distributed evaluation requires a matching device ShardLayout")
+ if len(layout.attrs) != len(mesh.layout.shape):
+ raise EvalError("distributed evaluation requires one attribute per mesh axis")
+ if not all(isinstance(attr, (Split, Partial, Broadcast)) for attr in layout.attrs):
+ raise EvalError("distributed evaluation does not support Dynamic ownership")
+ mapping = layout_axis_to_tensor_axis(layout.layout.shape, type_.shape)
+ targets = tuple(
+ mapping[attr.axis] if isinstance(attr, Split) else None for attr in layout.attrs
+ )
+ attrs = tuple(
+ Split(target) if target is not None else attr for target, attr in zip(targets, layout.attrs)
+ )
+ split_axes = [attr.axis for attr in attrs if isinstance(attr, Split)]
+ if len(set(split_axes)) != len(split_axes):
+ raise EvalError("distributed evaluation supports one mesh split per logical tensor axis")
+ for mesh_axis, attr in enumerate(layout.attrs):
+ if isinstance(attr, Split):
+ extent = layout.layout.shape[attr.axis]
+ preceding = layout.layout.shape[mapping.index(targets[mesh_axis]):attr.axis]
+ if any(size != 1 for size in preceding):
+ raise EvalError(
+ "distributed evaluation requires contiguous partitions at leading layout factors"
+ )
+ if isinstance(extent, int) and extent % mesh.layout.shape[mesh_axis]:
+ raise EvalError("split layout factor is not divisible by its mesh extent")
+ return attrs
+
+
+def _slices(type_, mesh, coord, bindings):
+ shape = _shape(type_, bindings)
+ attrs = _attrs(type_, mesh)
+ slices = []
+ for tensor_axis, size in enumerate(shape):
+ position, count = 0, 1
+ for mesh_axis, attr in enumerate(attrs):
+ if isinstance(attr, Split) and attr.axis == tensor_axis:
+ extent = mesh.layout.shape[mesh_axis]
+ position = position * extent + coord[mesh_axis]
+ count *= extent
+ if size % count:
+ raise EvalError(f"logical axis {tensor_axis} extent {size} is not divisible by {count}")
+ local = size // count
+ slices.append(slice(position * local, (position + 1) * local))
+ return tuple(slices)
+
+
+def distribute(value, type_, bindings):
+ """Bind an entry tensor to declared ownership without inventing partial values."""
+ if tuple(value.data.shape) != _shape(type_, bindings):
+ raise EvalError("logical input shape does not match its distributed declaration")
+ if not isinstance(type_.layout, ShardLayout):
+ return value
+ mesh = type_.layout.mesh
+ attrs = _attrs(type_, mesh)
+ if any(isinstance(attr, Partial) for attr in attrs):
+ raise EvalError("a logical input cannot supply Partial values")
+ return DistributedValue(
+ tuple(value.data[_slices(type_, mesh, coord, bindings)] for coord in coordinates(mesh)),
+ substitute_dims(type_, bindings),
+ mesh,
+ )
+
+
+def _local_type(type_, mesh, bindings):
+ slices = _slices(type_, mesh, coordinates(mesh)[0], bindings)
+ return replace(type_, shape=tuple(s.stop - s.start for s in slices), layout=None)
+
+
+def reconstruct(value):
+ """Assemble split outputs and check replicas; never complete a partial reduction."""
+ if isinstance(value, TupleValue):
+ return TupleValue(tuple(reconstruct(item) for item in value.elements))
+ if not isinstance(value, DistributedValue):
+ return value
+ if any(isinstance(attr, Partial) for attr in _attrs(value.type, value.mesh)):
+ raise EvalError(
+ "distributed output is Partial; an explicit reduction collective is required"
+ )
+ output = torch.empty(
+ value.type.shape, dtype=value.shards[0].dtype, device=value.shards[0].device
+ )
+ written = set()
+ for coord, shard in zip(coordinates(value.mesh), value.shards):
+ slices = _slices(value.type, value.mesh, coord, {})
+ if tuple(shard.shape) != tuple(s.stop - s.start for s in slices):
+ raise EvalError("distributed output does not match the declared shard shape")
+ key = tuple((s.start, s.stop) for s in slices)
+ if key in written:
+ try:
+ torch.testing.assert_close(output[slices], shard, rtol=0, atol=0, equal_nan=True)
+ except AssertionError as error:
+ raise EvalError("distributed output has inconsistent Broadcast replicas") from error
+ else:
+ output[slices] = shard
+ written.add(key)
+ return TensorValue(output, value.type)
+
+
+def _reshard(ctx, mesh, value):
+ source_attrs = _attrs(value.type, mesh)
+ dest_attrs = _attrs(ctx.result_type, mesh)
+ if any(isinstance(attr, Partial) for attr in dest_attrs) and dest_attrs != source_attrs:
+ raise EvalError("Reshard cannot create or complete Partial ownership; use a collective")
+ shards = (
+ value.shards
+ if isinstance(value, DistributedValue)
+ else (value.data,) * len(coordinates(mesh))
+ )
+ result = []
+ for coord, shard in zip(coordinates(mesh), shards):
+ source = _slices(value.type, mesh, coord, ctx.dim_bindings)
+ dest = _slices(ctx.result_type, mesh, coord, ctx.dim_bindings)
+ for before, after in zip(source_attrs, dest_attrs):
+ if isinstance(before, Partial) and before != after:
+ raise EvalError("Reshard cannot complete Partial ownership; use a collective")
+ if any(new.start < old.start or new.stop > old.stop for old, new in zip(source, dest)):
+ raise EvalError("Reshard requires cross-device data; use an explicit collective")
+ selection = tuple(
+ slice(new.start - old.start, new.stop - old.start) for old, new in zip(source, dest)
+ )
+ result.append(shard[selection])
+ return DistributedValue(tuple(result), substitute_dims(ctx.result_type, ctx.dim_bindings), mesh)
+
+
+def _collective(ctx, value):
+ mesh = value.mesh
+ axis = ctx.op.mesh_axis
+ coords = coordinates(mesh)
+ groups = {}
+ for rank, coord in enumerate(coords):
+ key = coord[:axis] + coord[axis + 1 :]
+ groups.setdefault(key, []).append(rank)
+ result_type = substitute_dims(ctx.result_type, ctx.dim_bindings)
+ _attrs(result_type, mesh)
+ source_attr = _attrs(value.type, mesh)[axis]
+ result = [None] * len(coords)
+ for ranks in groups.values():
+ shards = [value.shards[rank] for rank in ranks]
+ if isinstance(ctx.op, AllGather):
+ gathered = torch.cat(shards, dim=source_attr.axis)
+ for rank in ranks:
+ result[rank] = gathered
+ else:
+ combined = shards[0].clone()
+ for shard in shards[1:]:
+ if source_attr.reduction == "sum":
+ combined = combined + shard
+ elif source_attr.reduction == "max":
+ combined = torch.maximum(combined, shard)
+ else:
+ combined = torch.minimum(combined, shard)
+ if isinstance(ctx.op, ReduceScatter):
+ pieces = torch.tensor_split(combined, len(ranks), dim=ctx.op.tensor_axis)
+ for rank, piece in zip(ranks, pieces):
+ result[rank] = piece
+ else:
+ for rank in ranks:
+ result[rank] = combined
+ return DistributedValue(tuple(result), result_type, mesh)
+
+
+def evaluate_distributed_op(ctx, handler):
+ """Dispatch collectives collectively and supported local handlers per participant."""
+ if isinstance(ctx.op, TupleGetItem):
+ return handler(ctx)
+ distributed = [arg for arg in ctx.args if isinstance(arg, DistributedValue)]
+ mesh = ctx.mesh or (distributed[0].mesh if distributed else None)
+ if mesh is None:
+ result_layout = getattr(ctx.result_type, "layout", None)
+ if isinstance(result_layout, ShardLayout):
+ mesh = result_layout.mesh
+ if mesh is None:
+ return handler(ctx)
+ device_mesh_shape(mesh)
+ if any(arg.mesh != mesh for arg in distributed):
+ raise EvalError("distributed operands must use the current device mesh")
+ if isinstance(ctx.op, Reshard):
+ return _reshard(ctx, mesh, ctx.args[0])
+ if isinstance(ctx.op, (AllReduce, AllGather, ReduceScatter)):
+ if not isinstance(ctx.args[0], DistributedValue):
+ raise EvalError("collective requires distributed input")
+ return _collective(ctx, ctx.args[0])
+ if not isinstance(ctx.op, (Binary, Unary, MatMul, Cast, Transpose, Reduce)):
+ raise EvalError(
+ f"distributed local semantics are not supported for {type(ctx.op).__name__}"
+ )
+ if isinstance(ctx.op, Reduce):
+ reduced = {axis % len(ctx.args[0].type.shape) for axis in ctx.op.axes}
+ if any(
+ isinstance(attr, Split) and attr.axis in reduced
+ for attr in _attrs(ctx.args[0].type, mesh)
+ ):
+ raise EvalError(
+ "distributed Reduce over a Split axis needs explicit local reduction semantics"
+ )
+ result_type = _local_type(ctx.result_type, mesh, ctx.dim_bindings)
+ shards = []
+ for rank in range(len(coordinates(mesh))):
+ args = tuple(
+ TensorValue(arg.shards[rank], _local_type(arg.type, mesh, ctx.dim_bindings))
+ if isinstance(arg, DistributedValue)
+ else arg
+ for arg in ctx.args
+ )
+ result = handler(replace(ctx, args=args, result_type=result_type))
+ if not isinstance(result, TensorValue) or tuple(result.data.shape) != result_type.shape:
+ raise EvalError("local operation result does not match the declared shard shape")
+ shards.append(result.data)
+ return DistributedValue(tuple(shards), substitute_dims(ctx.result_type, ctx.dim_bindings), mesh)
diff --git a/src/tilefoundry/evaluator/interpreter.py b/src/tilefoundry/evaluator/interpreter.py
index 5e36be8c..9f56bf37 100644
--- a/src/tilefoundry/evaluator/interpreter.py
+++ b/src/tilefoundry/evaluator/interpreter.py
@@ -4,12 +4,19 @@
"""
from __future__ import annotations
+from dataclasses import replace
from typing import Any
import torch
from tilefoundry.evaluator.context import EvaluateContext
from tilefoundry.evaluator.dim import resolve_dim
+from tilefoundry.evaluator.distributed import (
+ DistributedValue,
+ distribute,
+ evaluate_distributed_op,
+ reconstruct,
+)
from tilefoundry.evaluator.registry import eval_registry
from tilefoundry.evaluator.value import (
EvalError,
@@ -24,6 +31,7 @@
from tilefoundry.ir.hir.function import Function
from tilefoundry.ir.hir.loop_region import LoopRegion
from tilefoundry.ir.hir.mesh_region import MeshRegion
+from tilefoundry.ir.hir.sharding.collective import device_mesh_shape
from tilefoundry.ir.types.dim import DimVar
from tilefoundry.ir.types.utils import types_compatible
from tilefoundry.ir.visitor import ExprVisitor
@@ -73,11 +81,12 @@ def _bind_dim_vars(params, values) -> dict[str, int]:
for p, v in zip(params, values):
shape = getattr(p.type, "shape", None)
data = getattr(v, "data", None)
- if shape is None or data is None:
+ data_shape = v.type.shape if isinstance(v, DistributedValue) else getattr(data, "shape", None)
+ if shape is None or data_shape is None:
continue
for axis, dim in enumerate(shape):
- if isinstance(dim, DimVar) and axis < len(data.shape):
- size = int(data.shape[axis])
+ if isinstance(dim, DimVar) and axis < len(data_shape):
+ size = int(data_shape[axis])
prev = binding.get(dim.name)
if prev is not None and prev != size:
raise EvalError(
@@ -122,6 +131,11 @@ def visit_MeshRegion(self, region: MeshRegion, ctx: EvaluateContext) -> Value:
f"evaluator: MeshRegion expects {len(region.params)} args, got {len(args)}"
)
memo = {id(param): (param, value) for param, value in zip(region.params, args)}
+ if ctx.distributed:
+ device_mesh_shape(region.mesh)
+ if ctx.mesh is not None and ctx.mesh != region.mesh:
+ raise EvalError("distributed evaluation does not support nested different meshes")
+ ctx = replace(ctx, mesh=region.mesh)
return EvaluatorVisitor(memo=memo).visit(region.body, ctx)
def visit_leaf_Var(self, var: Var, _operands, ctx: EvaluateContext) -> Value:
@@ -146,7 +160,10 @@ def visit_leaf_Call(self, call: Call, args, ctx: EvaluateContext) -> Value:
handler = eval_registry.lookup(type(target))
if handler is None:
raise EvalError(f"no @register_eval handler for {type(target).__name__}")
- return handler(ctx.for_op(target, args, call.type))
+ op_context = ctx.for_op(target, args, call.type)
+ if ctx.distributed:
+ return evaluate_distributed_op(op_context, handler)
+ return handler(op_context)
except Exception as error:
raise EvalError(f"evaluator: {describe_expr(call)}: {error}") from error
@@ -181,7 +198,15 @@ def _call_function(
loaded_module=child if child is not None else ctx.loaded_module,
device=ctx.device,
dim_bindings=_bind_dim_vars(target.params, args),
+ distributed=ctx.distributed,
)
+ if ctx.distributed:
+ args = [
+ distribute(arg, param.type, function_context.dim_bindings)
+ if isinstance(arg, TensorValue)
+ else arg
+ for param, arg in zip(target.params, args)
+ ]
memo = {id(param): (param, arg) for param, arg in zip(target.params, args)}
return EvaluatorVisitor(memo=memo).visit(target.body, function_context)
@@ -304,7 +329,9 @@ def _child_constant(loaded_module, callee: Function, param, device: str) -> Tens
return TensorValue(data=value, type=param.type)
-def _run_selected(loaded_module, fn: Function, *activations, device: str | None):
+def _run_selected(
+ loaded_module, fn: Function, *activations, device: str | None, distributed: bool = False
+):
"""Run one loaded function, reading constants lazily at first use.
*device* is where the activations already are; a weight somewhere else is
@@ -323,7 +350,7 @@ def _run_selected(loaded_module, fn: Function, *activations, device: str | None)
else next(supplied)
for param in fn.params
]
- return _run_bound(fn, args, device=device, reading=loaded_module)
+ return _run_bound(fn, args, device=device, reading=loaded_module, distributed=distributed)
def _unwrap(value: Value) -> Any:
@@ -347,10 +374,13 @@ def _select_variant(callee: Function, arg_values) -> Function:
if loc is None:
continue
pi, axis = loc
- data = getattr(arg_values[pi], "data", None)
- if data is None or axis >= len(data.shape):
+ value = arg_values[pi]
+ data_shape = value.type.shape if isinstance(value, DistributedValue) else getattr(
+ getattr(value, "data", None), "shape", None
+ )
+ if data_shape is None or axis >= len(data_shape):
continue
- if pat.match(int(data.shape[axis])):
+ if pat.match(int(data_shape[axis])):
matches.append(v)
if len(matches) != 1:
raise EvalError(
@@ -383,7 +413,9 @@ def _selected_body(fn: Function, args) -> Function:
return _select_variant(fn, _bound_values(fn, args))
-def _run_bound(fn: Function, args, *, device: str | None = None, reading=None):
+def _run_bound(
+ fn: Function, args, *, device: str | None = None, reading=None, distributed: bool = False
+):
"""Evaluate *fn* over fully bound *args*, with *reading* in hand.
The entry a resource reading runs through: every child call reached from
@@ -392,21 +424,23 @@ def _run_bound(fn: Function, args, *, device: str | None = None, reading=None):
"""
values = _bound_values(fn, args)
target = _select_variant(fn, values) if fn.variants else fn
- memo = {id(param): (param, value) for param, value in zip(target.params, values)}
dim_env = _bind_dim_vars(target.params, values)
- return _unwrap(
- EvaluatorVisitor(memo=memo).visit(
- target.body,
- EvaluateContext(
- loaded_module=reading,
- device=device,
- dim_bindings=dim_env,
- ),
- )
+ if distributed:
+ values = [distribute(value, param.type, dim_env) for param, value in zip(target.params, values)]
+ memo = {id(param): (param, value) for param, value in zip(target.params, values)}
+ result = EvaluatorVisitor(memo=memo).visit(
+ target.body,
+ EvaluateContext(
+ loaded_module=reading,
+ device=device,
+ dim_bindings=dim_env,
+ distributed=distributed,
+ ),
)
+ return _unwrap(reconstruct(result) if distributed else result)
-def evaluate(target, *inputs):
+def evaluate(target, *inputs, distributed: bool = False):
"""Evaluate a HIR ``Function`` or a loaded module and return torch value(s).
*target* answers which function to run and which reading, if any, supplies
@@ -416,10 +450,10 @@ def evaluate(target, *inputs):
fn, reading = target.evaluation_target()
device = _device_of(inputs)
if reading is not None:
- return _run_selected(reading, fn, *inputs, device=device)
+ return _run_selected(reading, fn, *inputs, device=device, distributed=distributed)
if len(inputs) != len(fn.params):
raise EvalError(
f"evaluator: {fn.name!r} expects {len(fn.params)} inputs, got {len(inputs)}"
)
- return _run_bound(fn, inputs, device=device)
+ return _run_bound(fn, inputs, device=device, distributed=distributed)
diff --git a/src/tilefoundry/ir/hir/sharding/collective.py b/src/tilefoundry/ir/hir/sharding/collective.py
new file mode 100644
index 00000000..2908f134
--- /dev/null
+++ b/src/tilefoundry/ir/hir/sharding/collective.py
@@ -0,0 +1,132 @@
+"""Explicit communication over one axis of a device mesh."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from math import prod
+
+from tilefoundry.evaluator.registry import register_eval
+from tilefoundry.evaluator.value import EvalError
+from tilefoundry.ir.core import Op
+from tilefoundry.ir.core.param_def import ParamDef
+from tilefoundry.ir.core.pattern import Tensor
+from tilefoundry.ir.core.register import register_op
+from tilefoundry.ir.types.shard import Layout, Mesh, c_order_strides, canonical_shard_layout
+from tilefoundry.ir.types.shard.scope_match import covered_by_scope
+from tilefoundry.ir.types.shard.shard_layout import (
+ Broadcast,
+ Partial,
+ ShardLayout,
+ Split,
+ layout_axis_to_tensor_axis,
+ split_target_axes,
+)
+from tilefoundry.visitor_registry import register_typeinfer
+from tilefoundry.visitor_registry.access_relation import (
+ identity_relations,
+ register_access_relation,
+)
+
+
+def device_mesh_shape(mesh: Mesh) -> tuple[int, ...]:
+ """Validate the concrete, single-level device mesh supported by collectives."""
+ if len(mesh.topologies) != 1 or mesh.topologies[0].name != "gpu":
+ raise ValueError("collectives require a single-level gpu mesh")
+ if not isinstance(mesh.layout, Layout):
+ raise ValueError("collectives require a concrete Layout mesh")
+ shape = mesh.layout.shape
+ if not shape or any(type(size) is not int or size <= 0 for size in shape):
+ raise ValueError("collectives require positive concrete mesh extents")
+ if mesh.layout.strides != c_order_strides(shape):
+ raise ValueError("collectives require a contiguous mesh in coordinate order")
+ capacity = mesh.topologies[0].size
+ if type(capacity) is not int or prod(shape) > capacity:
+ raise ValueError("collective mesh exceeds its concrete topology extent")
+ return shape
+
+
+@register_op
+class AllReduce(Op):
+ """Complete a partial reduction and replicate it along one mesh axis."""
+
+ x = ParamDef(kind="input", pattern=Tensor)
+ mesh_axis = ParamDef(kind="attribute", annotation=int)
+
+
+@register_op
+class AllGather(Op):
+ """Gather equal shards in coordinate order along one mesh axis."""
+
+ x = ParamDef(kind="input", pattern=Tensor)
+ mesh_axis = ParamDef(kind="attribute", annotation=int)
+
+
+@register_op
+class ReduceScatter(Op):
+ """Complete a partial reduction and equally partition one logical axis."""
+
+ x = ParamDef(kind="input", pattern=Tensor)
+ mesh_axis = ParamDef(kind="attribute", annotation=int)
+ tensor_axis = ParamDef(kind="attribute", annotation=int)
+
+
+def _collective_type(call, ctx):
+ source = ctx.type_of(call.args[0])
+ layout = source.layout
+ if not isinstance(layout, ShardLayout):
+ ctx.error(call, "collective input must have a ShardLayout")
+ try:
+ shape = device_mesh_shape(layout.mesh)
+ except ValueError as error:
+ ctx.error(call, str(error))
+ if ctx.current_mesh is None or not covered_by_scope(layout.mesh, ctx.current_mesh):
+ ctx.error(call, "collective mesh must be bound by the current mesh scope")
+ axis = call.target.mesh_axis
+ if type(axis) is not int or not 0 <= axis < len(shape):
+ ctx.error(call, "mesh_axis must index the input mesh")
+ if len(layout.attrs) != len(shape):
+ ctx.error(call, "collective input needs one shard attribute per mesh axis")
+ if not all(isinstance(attr, (Broadcast, Split, Partial)) for attr in layout.attrs):
+ ctx.error(call, "collectives support only Broadcast, Split and Partial ownership")
+ attrs = list(layout.attrs)
+ targets = split_target_axes(layout, source.shape)
+ mapping = layout_axis_to_tensor_axis(layout.layout.shape, source.shape)
+ for attr, target in zip(attrs, targets):
+ if isinstance(attr, Split):
+ preceding = layout.layout.shape[mapping.index(target):attr.axis]
+ if any(size != 1 for size in preceding):
+ ctx.error(call, "collectives require contiguous partitions at leading layout factors")
+ attrs = [Split(target) if target is not None else attr for target, attr in zip(targets, attrs)]
+ split_axes = [attr.axis for attr in attrs if isinstance(attr, Split)]
+ if len(set(split_axes)) != len(split_axes):
+ ctx.error(call, "collectives support one mesh split per logical tensor axis")
+ selected = attrs[axis]
+ if isinstance(call.target, AllGather):
+ if not isinstance(selected, Split):
+ ctx.error(call, "AllGather requires Split on mesh_axis")
+ elif not isinstance(selected, Partial) or selected.reduction not in ("sum", "max", "min"):
+ ctx.error(call, "reduction collective requires Partial(sum, max or min) on mesh_axis")
+ if isinstance(call.target, ReduceScatter):
+ tensor_axis = call.target.tensor_axis
+ if type(tensor_axis) is not int or not 0 <= tensor_axis < len(source.shape):
+ ctx.error(call, "tensor_axis must index the logical tensor shape")
+ if tensor_axis in split_axes:
+ ctx.error(call, "ReduceScatter tensor_axis is already split by another mesh axis")
+ attrs[axis] = Split(tensor_axis)
+ else:
+ attrs[axis] = Broadcast()
+ try:
+ result = canonical_shard_layout(source.shape, layout.mesh, tuple(attrs))
+ except ValueError as error:
+ ctx.error(call, str(error))
+ return replace(source, layout=result)
+
+
+def _eval_collective(ctx):
+ raise EvalError("device collectives require distributed evaluation")
+
+
+for _op in (AllReduce, AllGather, ReduceScatter):
+ register_typeinfer(_op)(_collective_type)
+ register_eval(_op)(_eval_collective)
+ register_access_relation(_op)(identity_relations(1))
diff --git a/src/tilefoundry/ir/types/shard/shard_layout.py b/src/tilefoundry/ir/types/shard/shard_layout.py
index 67d01fb9..c605ede7 100644
--- a/src/tilefoundry/ir/types/shard/shard_layout.py
+++ b/src/tilefoundry/ir/types/shard/shard_layout.py
@@ -215,7 +215,9 @@ def layout_axis_to_tensor_axis(layout_shape: tuple, tensor_shape: tuple) -> list
layout_idx += 1
continue
running = 1
- while layout_idx < len(layout_shape) and running < t_dim_int:
+ while layout_idx < len(layout_shape) and (
+ running < t_dim_int or (t_dim_int == 0 and running != 0)
+ ):
sh = static_dim_value(layout_shape[layout_idx])
running *= 1 if sh is None else sh
result.append(t_axis)
diff --git a/tests/evaluator/test_distributed.py b/tests/evaluator/test_distributed.py
new file mode 100644
index 00000000..3efca3c9
--- /dev/null
+++ b/tests/evaluator/test_distributed.py
@@ -0,0 +1,181 @@
+"""Numerical contracts of distributed HIR through public evaluator entry points."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tests._source import import_dsl
+from tests.fixtures.distributed.projection import ReduceScatterProjection, Reference, TensorParallel
+from tilefoundry.evaluator import EvalError, evaluate
+from tilefoundry.ir.core import VerifyError
+from tilefoundry.runtime import DictResource
+
+
+@pytest.mark.parametrize("candidate", [TensorParallel, ReduceScatterProjection])
+def test_projection_and_repeated_state(candidate):
+ generator = torch.Generator().manual_seed(81)
+ weight = torch.randn(8, 6, generator=generator)
+ reference_state = torch.randn(4, 6, generator=generator)
+ candidate_state = reference_state.clone()
+ resource = DictResource({"w": weight})
+ for _ in range(3):
+ x = torch.randn(4, 8, generator=generator)
+ expected = x @ weight + reference_state
+ reference_output, reference_state = evaluate(Reference.load(resource), x, reference_state)
+ output, candidate_state = evaluate(
+ candidate.load(resource), x, candidate_state, distributed=True
+ )
+ torch.testing.assert_close(reference_output, -expected)
+ torch.testing.assert_close(output, -expected)
+ torch.testing.assert_close(candidate_state, expected)
+
+
+_PROJECTION = """
+from tilefoundry import module, func
+from tilefoundry.dsl import Mesh, Tensor, Topology, tf
+@module(entry="run", topologies=(Topology("gpu", 4),))
+class Projection:
+ @func
+ def run(x: Tensor[(4, 8), "f32"], w: Tensor[(8, 6), "f32"]):
+ with Mesh(("gpu",), (2, 2), names=("dp", "tp")) as devices:
+ a = tf.reshard(x, (4 @ devices.dp, 8 @ devices.tp), "gmem")
+ b = tf.reshard(w, (8 @ devices.tp, 6), "gmem")
+ partial = tf.matmul(a, b)
+ complete = tf.allreduce(partial, mesh_axis=1)
+ return tf.allgather(complete, mesh_axis=0)
+"""
+
+
+def test_multiaxis_groups_gather_in_coordinate_order():
+ module = import_dsl(_PROJECTION, "Projection")
+ x = torch.arange(32, dtype=torch.float32).reshape(4, 8)
+ w = torch.arange(48, dtype=torch.float32).reshape(8, 6)
+ torch.testing.assert_close(evaluate(module.entry_function(), x, w, distributed=True), x @ w)
+
+
+@pytest.mark.parametrize("batch", [0, 2, 3, 6])
+def test_dynamic_partitions_check_divisibility_and_empty_payloads(batch):
+ source = (
+ _PROJECTION.replace(
+ "from tilefoundry import module, func",
+ "from tilefoundry import module, func\n"
+ 'from tilefoundry.ir.types.dim import DimVar\nN = DimVar("batch", 0, 8)',
+ )
+ .replace("Tensor[(4, 8)", "Tensor[(N, 8)")
+ .replace("(4 @ devices.dp,", "(N @ devices.dp,")
+ )
+ module = import_dsl(source, "Projection")
+ x = torch.arange(batch * 8, dtype=torch.float32).reshape(batch, 8)
+ w = torch.arange(48, dtype=torch.float32).reshape(8, 6)
+ if batch % 2:
+ with pytest.raises(EvalError, match="not divisible"):
+ evaluate(module.entry_function(), x, w, distributed=True)
+ else:
+ torch.testing.assert_close(evaluate(module.entry_function(), x, w, distributed=True), x @ w)
+
+
+@pytest.mark.parametrize(
+ "replacement, message",
+ [
+ ("return partial", "output is Partial"),
+ ("return tf.reshard(partial, (4 @ devices.dp, 6), 'gmem')", "cannot complete Partial"),
+ ("return tf.reshard(complete, (4, 6), 'gmem')", "cross-device data"),
+ ],
+)
+def test_missing_communication_is_not_repaired(replacement, message):
+ source = _PROJECTION.replace("return tf.allgather(complete, mesh_axis=0)", replacement)
+ module = import_dsl(source, "Projection")
+ with pytest.raises(EvalError, match=message):
+ evaluate(module.entry_function(), torch.ones(4, 8), torch.ones(8, 6), distributed=True)
+
+
+@pytest.mark.parametrize(
+ "source, message",
+ [
+ (
+ _PROJECTION.replace(
+ "allreduce(partial, mesh_axis=1)", "allreduce(partial, mesh_axis=0)"
+ ),
+ "requires Partial",
+ ),
+ (
+ _PROJECTION.replace(
+ "allgather(complete, mesh_axis=0)", "allgather(complete, mesh_axis=1)"
+ ),
+ "requires Split",
+ ),
+ (
+ _PROJECTION.replace(
+ "allreduce(partial, mesh_axis=1)", "allreduce(partial, mesh_axis=3)"
+ ),
+ "mesh_axis",
+ ),
+ ],
+)
+def test_illegal_collective_contracts_are_rejected(source, message):
+ with pytest.raises(VerifyError, match=message):
+ import_dsl(source, "Projection")
+
+
+def test_collectives_require_distributed_execution():
+ module = import_dsl(_PROJECTION, "Projection")
+ with pytest.raises(EvalError, match="require distributed evaluation"):
+ evaluate(module.entry_function(), torch.ones(4, 8), torch.ones(8, 6))
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+def test_local_composition_preserves_transposed_ownership(axis):
+ source = """
+from tilefoundry import module, func
+from tilefoundry.dsl import Mesh, Tensor, Topology, tf
+@module(entry="run", topologies=(Topology("gpu", 2),))
+class LocalComposition:
+ @func
+ def run(x: Tensor[(4, 6), "f32"]):
+ with Mesh(("gpu",), (2,), names=("tp",)) as devices:
+ local = tf.reshard(x, (4 @ devices.tp, 6), "gmem")
+ transposed = tf.transpose(local, perm=(1, 0))
+ rounded = tf.cast(transposed, dtype="bf16")
+ return tf.reduce(rounded, axes=(AXIS,), keepdim=False, kind="sum")
+""".replace("AXIS", str(axis))
+ module = import_dsl(source, "LocalComposition")
+ x = torch.randn(4, 6, generator=torch.Generator().manual_seed(42))
+ if axis == 1:
+ with pytest.raises(EvalError, match="Reduce over a Split axis"):
+ evaluate(module.entry_function(), x, distributed=True)
+ else:
+ expected = x.T.bfloat16().sum(dim=0)
+ torch.testing.assert_close(evaluate(module.entry_function(), x, distributed=True), expected)
+
+
+def test_collectives_survive_function_and_loop_boundaries():
+ module = import_dsl(
+ """
+from __future__ import annotations
+from tilefoundry import module, func
+from tilefoundry.dsl import Mesh, Tensor, Topology, tf
+@module(entry="run", topologies=(Topology("gpu", 2),))
+class Repeated:
+ @func(mesh=Mesh(("gpu",), (2,), names=("tp",)))
+ def project(x: Tensor[(4, 8 @ mesh.tp), "f32"], w: Tensor[(8 @ mesh.tp, 6), "f32"]):
+ return tf.allreduce(tf.matmul(x, w), mesh_axis=0)
+
+ @func
+ def run(x: Tensor[(4, 8), "f32"], w: Tensor[(8, 6), "f32"], state: Tensor[(4, 6), "f32"]):
+ with Mesh(("gpu",), (2,), names=("tp",)) as devices:
+ a = tf.reshard(x, (4, 8 @ devices.tp), "gmem")
+ b = tf.reshard(w, (8 @ devices.tp, 6), "gmem")
+ current = tf.reshard(state, (4, 6), "gmem")
+ for i in range(3):
+ current = tf.add(current, project(a, b))
+ return current
+""",
+ "Repeated",
+ )
+ generator = torch.Generator().manual_seed(14)
+ x = torch.randn(4, 8, generator=generator)
+ w = torch.randn(8, 6, generator=generator)
+ state = torch.randn(4, 6, generator=generator)
+ output = evaluate(module.entry_function(), x, w, state, distributed=True)
+ torch.testing.assert_close(output, state + 3 * (x @ w))
diff --git a/tests/fixtures/distributed/__init__.py b/tests/fixtures/distributed/__init__.py
new file mode 100644
index 00000000..fdd4db4d
--- /dev/null
+++ b/tests/fixtures/distributed/__init__.py
@@ -0,0 +1 @@
+"""Portable examples of explicit device communication."""
diff --git a/tests/fixtures/distributed/projection.py b/tests/fixtures/distributed/projection.py
new file mode 100644
index 00000000..748e8bc9
--- /dev/null
+++ b/tests/fixtures/distributed/projection.py
@@ -0,0 +1,40 @@
+"""Reference and distributed projections with explicit recurrent state."""
+
+from tilefoundry import func, module
+from tilefoundry.dsl import ConstTensor, Mesh, Tensor, Topology, tf
+
+
+@module(entry="step")
+class Reference:
+ @func
+ def step(x: Tensor[(4, 8), "f32"], w: ConstTensor[(8, 6), "f32"], state: Tensor[(4, 6), "f32"]):
+ updated = tf.add(tf.matmul(x, w), state)
+ return tf.neg(updated), updated
+
+
+@module(entry="step", topologies=(Topology("gpu", 2),))
+class TensorParallel:
+ @func
+ def step(x: Tensor[(4, 8), "f32"], w: ConstTensor[(8, 6), "f32"], state: Tensor[(4, 6), "f32"]):
+ with Mesh(("gpu",), (2,), names=("tp",)) as devices:
+ local_x = tf.reshard(x, (4, 8 @ devices.tp), "gmem")
+ local_w = tf.reshard(w, (8 @ devices.tp, 6), "gmem")
+ partial = tf.matmul(local_x, local_w)
+ complete = tf.allreduce(partial, mesh_axis=0)
+ replicated_state = tf.reshard(state, (4, 6), "gmem")
+ updated = tf.add(complete, replicated_state)
+ return tf.neg(updated), updated
+
+
+@module(entry="step", topologies=(Topology("gpu", 2),))
+class ReduceScatterProjection:
+ @func
+ def step(x: Tensor[(4, 8), "f32"], w: ConstTensor[(8, 6), "f32"], state: Tensor[(4, 6), "f32"]):
+ with Mesh(("gpu",), (2,), names=("tp",)) as devices:
+ local_x = tf.reshard(x, (4, 8 @ devices.tp), "gmem")
+ local_w = tf.reshard(w, (8 @ devices.tp, 6), "gmem")
+ partial = tf.matmul(local_x, local_w)
+ complete = tf.reducescatter(partial, mesh_axis=0, tensor_axis=0)
+ local_state = tf.reshard(state, (4 @ devices.tp, 6), "gmem")
+ updated = tf.add(complete, local_state)
+ return tf.neg(updated), updated
From 8ce779979b1cd5fcbcc384638a5a674b4d38808f Mon Sep 17 00:00:00 2001
From: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
Date: Wed, 23 Sep 2026 16:48:26 +0800
Subject: [PATCH 2/2] feat(cli): compare distributed HIR with explicit
references [M1]
---
docs/spec/cli.md | 19 ++++-
docs/tutorial/distributed-check.ipynb | 72 +++++++++++++++++
docs/tutorial/distributed-check.md | 111 ++++++++++++++++++++++++++
docs/tutorial/index.md | 4 +
src/tilefoundry/cli/check.py | 67 +++++++++++++++-
src/tilefoundry/cli/tutorial.py | 2 +-
tests/cli/test_cli_check.py | 49 ++++++++++++
7 files changed, 318 insertions(+), 6 deletions(-)
create mode 100644 docs/tutorial/distributed-check.ipynb
create mode 100644 docs/tutorial/distributed-check.md
diff --git a/docs/spec/cli.md b/docs/spec/cli.md
index 6ccc9991..f120f377 100644
--- a/docs/spec/cli.md
+++ b/docs/spec/cli.md
@@ -168,6 +168,20 @@ bounds the caller stated.
authored Module the implementation stands for. With no reference at all, only
a predicate that judges the candidate alone is admissible; every two-sided
predicate MUST be refused, because there is nothing to compare against.
+ - `--reference SOURCE` MAY select an explicit reference HIR using the same
+ source/selector syntax as the candidate. Both selections MUST be HIR modules
+ or functions; runtime twins and `--expected` MUST NOT accompany this mode.
+ Activation parameters MUST match in declaration order by logical shape and
+ dtype. The selected module trees MUST declare matching weight paths, shapes
+ and dtypes. Placement annotations MAY differ. Both programs receive the
+ same activations and the same weight resource scoped at the candidate's
+ selected module; relative child/weight paths MUST retain their meaning.
+ - `--distributed` enables rank simulation for the HIR candidate according to
+ [evaluator §7](./evaluator.md#7-distributed-evaluation). The explicit reference
+ uses ordinary logical evaluation. The flag MUST be refused for a runtime
+ twin. Existing expected-output and candidate-only predicates remain usable.
+ Text and JSON reports MUST identify the explicit reference source and the
+ candidate evaluation mode (`local` or `distributed`).
- Each output MUST report the norm of its reference. Near zero, a relative
measure divides by nothing, so the report MUST state what it measured instead
rather than a number with no scale to read it against.
@@ -235,9 +249,12 @@ granularity.
every source file's leading docstring without importing or executing it. An
unknown family MUST name the available families. Checkout and installed
lookups MUST report the same shipped families and files.
- - Its pages are `index`, `migrate`, `optimize`, and `showcase`; the first three
+ - Its pages are `index`, `migrate`, `optimize`, `showcase`, and
+ `distributed-check`; the first three
are the workflow and `showcase` is one kernel taken through six analyze-driven
stages, exercising the authoring surface the other pages touch in part.
+ `distributed-check` compares a device-parallel HIR with its reference using
+ simulated ranks and explicit collectives.
Causal-LM decode sources are listed through `tutorial orchestrator`. A bare
`tutorial` MUST print the `index` page followed by its own help, which names
the pages a reader may ask for and `orchestrator`. `index` is that overview's
diff --git a/docs/tutorial/distributed-check.ipynb b/docs/tutorial/distributed-check.ipynb
new file mode 100644
index 00000000..937015e7
--- /dev/null
+++ b/docs/tutorial/distributed-check.ipynb
@@ -0,0 +1,72 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "# Check a distributed projection against its reference\n\nUse `check --reference SOURCE --distributed` to evaluate a transformed HIR\nwith simulated device ranks and compare its logical outputs against another\nHIR. The simulation runs on the torch device supplied by `--device`, so the\nexample below runs on a CPU.\n\nSave this page as `distributed-check.md`, then extract its program:"
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "tilefoundry": {
+ "cell_type": "bash",
+ "command": "extract-source"
+ }
+ },
+ "outputs": [],
+ "source": "%%bash\nset -euo pipefail\nawk -v tag=\"\" '\n $0 == tag { block=1; next }\n block && /^```python$/ { in_python=1; next }\n in_python && /^```$/ { in_python=0; block=0; next }\n in_python { print }\n' distributed-check.md > projection.py"
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {
+ "tilefoundry": {
+ "source": "projection.py"
+ }
+ },
+ "outputs": [],
+ "source": "# example\nfrom tilefoundry import func, module\nfrom tilefoundry.dsl import ConstTensor, Mesh, Tensor, Topology, tf\n\n\n@module(entry=\"step\")\nclass Reference:\n @func\n def step(\n x: Tensor[(4, 8), \"f32\"],\n w: ConstTensor[(8, 6), \"f32\"],\n state: Tensor[(4, 6), \"f32\"],\n ):\n updated = tf.add(tf.matmul(x, w), state)\n return tf.neg(updated), updated\n\n\n@module(entry=\"step\", topologies=(Topology(\"gpu\", 2),))\nclass TensorParallel:\n @func\n def step(\n x: Tensor[(4, 8), \"f32\"],\n w: ConstTensor[(8, 6), \"f32\"],\n state: Tensor[(4, 6), \"f32\"],\n ):\n with Mesh((\"gpu\",), (2,), names=(\"tp\",)) as devices:\n a = tf.reshard(x, (4, 8 @ devices.tp), \"gmem\")\n b = tf.reshard(w, (8 @ devices.tp, 6), \"gmem\")\n partial = tf.matmul(a, b)\n complete = tf.allreduce(partial, mesh_axis=0)\n replicated_state = tf.reshard(state, (4, 6), \"gmem\")\n updated = tf.add(complete, replicated_state)\n return tf.neg(updated), updated"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "Each rank computes half the contraction. Its result is a partial sum;\n`allreduce` completes that sum before adding state. The second output explicitly\ncarries the state for the next invocation."
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {
+ "tilefoundry": {
+ "cell_type": "bash"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": "projection.py:TensorParallel\n reference: projection.py:Reference\n evaluation: distributed\n inputs: random (seed 0); activations actual f32, f32 (declared f32, f32)\n\n output[0] f32[4,6] ref_norm 17.3306\n allclose(atol=1e-05 rtol=1e-05) max_violation 0 PASS\n output[1] f32[4,6] ref_norm 17.3306\n allclose(atol=1e-05 rtol=1e-05) max_violation 0 PASS\n\nPASS\n"
+ }
+ ],
+ "source": "%%bash\nset -euo pipefail\ntilefoundry check projection.py:TensorParallel \\\n --reference projection.py:Reference --distributed \\\n --inputs random --weights random --device cpu \\\n --out 'output[0]' --fn allclose --atol 1e-5 --rtol 1e-5 \\\n --out 'output[1]' --fn allclose --atol 1e-5 --rtol 1e-5"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "Both programs receive the same logical activations and weight resource. The\nchosen tolerances are explicit for this small floating-point example; choose\nthem for the numerical policy of the actual program. Add `--json report.json`\nto retain a machine-readable result identifying the candidate, reference and\nevaluation mode.\n\nReturning the partial sum raises an error. Replacing `allreduce` with a\n`Reshard` to broadcast ownership also raises an error: a layout annotation\ncannot silently perform cross-device communication. `allgather` assembles an\nexisting split, while `reducescatter` completes a partial reduction and leaves\nits result split. The evaluator reconstructs split outputs for comparison.\n\nFor a stateful check, call `evaluate(candidate.load(resource), x, state,\ndistributed=True)` repeatedly and pass the returned state into the next call.\nKeep an independent reference state and compare both the visible output and\nthe next state at every step. The CLI checks one invocation per dimension point;\nit does not infer a state feedback loop.\n\nThe [evaluator contract](../spec/evaluator.md#7-distributed-evaluation) lists\nsupported local operations and partition shapes. This checks collective value\nsemantics, not a communication backend or performance. `analyze` currently\nreports device collectives as unsupported because their communication cost\nmodels are not registered.\n\n\nUse `tilefoundry check --help` for comparison predicates and\n`tilefoundry spec evaluator` for the evaluation contract."
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.12"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/tutorial/distributed-check.md b/docs/tutorial/distributed-check.md
new file mode 100644
index 00000000..5a3f2f42
--- /dev/null
+++ b/docs/tutorial/distributed-check.md
@@ -0,0 +1,111 @@
+# Check a distributed projection against its reference
+
+Use `check --reference SOURCE --distributed` to evaluate a transformed HIR
+with simulated device ranks and compare its logical outputs against another
+HIR. The simulation runs on the torch device supplied by `--device`, so the
+example below runs on a CPU.
+
+Save this page as `distributed-check.md`, then extract its program:
+
+```bash
+set -euo pipefail
+awk -v tag="" '
+ $0 == tag { block=1; next }
+ block && /^```python$/ { in_python=1; next }
+ in_python && /^```$/ { in_python=0; block=0; next }
+ in_python { print }
+' distributed-check.md > projection.py
+```
+
+
+
+```python
+# example
+from tilefoundry import func, module
+from tilefoundry.dsl import ConstTensor, Mesh, Tensor, Topology, tf
+
+
+@module(entry="step")
+class Reference:
+ @func
+ def step(
+ x: Tensor[(4, 8), "f32"],
+ w: ConstTensor[(8, 6), "f32"],
+ state: Tensor[(4, 6), "f32"],
+ ):
+ updated = tf.add(tf.matmul(x, w), state)
+ return tf.neg(updated), updated
+
+
+@module(entry="step", topologies=(Topology("gpu", 2),))
+class TensorParallel:
+ @func
+ def step(
+ x: Tensor[(4, 8), "f32"],
+ w: ConstTensor[(8, 6), "f32"],
+ state: Tensor[(4, 6), "f32"],
+ ):
+ with Mesh(("gpu",), (2,), names=("tp",)) as devices:
+ a = tf.reshard(x, (4, 8 @ devices.tp), "gmem")
+ b = tf.reshard(w, (8 @ devices.tp, 6), "gmem")
+ partial = tf.matmul(a, b)
+ complete = tf.allreduce(partial, mesh_axis=0)
+ replicated_state = tf.reshard(state, (4, 6), "gmem")
+ updated = tf.add(complete, replicated_state)
+ return tf.neg(updated), updated
+```
+
+Each rank computes half the contraction. Its result is a partial sum;
+`allreduce` completes that sum before adding state. The second output explicitly
+carries the state for the next invocation.
+
+```bash
+set -euo pipefail
+tilefoundry check projection.py:TensorParallel \
+ --reference projection.py:Reference --distributed \
+ --inputs random --weights random --device cpu \
+ --out 'output[0]' --fn allclose --atol 1e-5 --rtol 1e-5 \
+ --out 'output[1]' --fn allclose --atol 1e-5 --rtol 1e-5
+```
+
+```text
+projection.py:TensorParallel
+ reference: projection.py:Reference
+ evaluation: distributed
+ inputs: random (seed 0); activations actual f32, f32 (declared f32, f32)
+
+ output[0] f32[4,6] ref_norm 17.3306
+ allclose(atol=1e-05 rtol=1e-05) max_violation 0 PASS
+ output[1] f32[4,6] ref_norm 17.3306
+ allclose(atol=1e-05 rtol=1e-05) max_violation 0 PASS
+
+PASS
+```
+
+Both programs receive the same logical activations and weight resource. The
+chosen tolerances are explicit for this small floating-point example; choose
+them for the numerical policy of the actual program. Add `--json report.json`
+to retain a machine-readable result identifying the candidate, reference and
+evaluation mode.
+
+Returning the partial sum raises an error. Replacing `allreduce` with a
+`Reshard` to broadcast ownership also raises an error: a layout annotation
+cannot silently perform cross-device communication. `allgather` assembles an
+existing split, while `reducescatter` completes a partial reduction and leaves
+its result split. The evaluator reconstructs split outputs for comparison.
+
+For a stateful check, call `evaluate(candidate.load(resource), x, state,
+distributed=True)` repeatedly and pass the returned state into the next call.
+Keep an independent reference state and compare both the visible output and
+the next state at every step. The CLI checks one invocation per dimension point;
+it does not infer a state feedback loop.
+
+The [evaluator contract](../spec/evaluator.md#7-distributed-evaluation) lists
+supported local operations and partition shapes. This checks collective value
+semantics, not a communication backend or performance. `analyze` currently
+reports device collectives as unsupported because their communication cost
+models are not registered.
+
+
+Use `tilefoundry check --help` for comparison predicates and
+`tilefoundry spec evaluator` for the evaluation contract.
diff --git a/docs/tutorial/index.md b/docs/tutorial/index.md
index a69e35c4..db146900 100644
--- a/docs/tutorial/index.md
+++ b/docs/tutorial/index.md
@@ -4,6 +4,10 @@ TileFoundry is source to source: the reference is source code, the fast
implementation is source code, and either can be pointed at any command.
`check` says whether two of them agree; `analyze` says what one costs.
+For device-parallel HIR, [distributed checking](./distributed-check.md) shows
+how to compare explicit collectives and state outputs against a reference using
+`check --reference ... --distributed`.
+
```text
step one — describe it, until it agrees
diff --git a/src/tilefoundry/cli/check.py b/src/tilefoundry/cli/check.py
index e9e0f194..3a3a2558 100644
--- a/src/tilefoundry/cli/check.py
+++ b/src/tilefoundry/cli/check.py
@@ -63,9 +63,16 @@ def add_arguments(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"source", metavar="SOURCE", help="FILE.py:Selector — a Module, a leaf, or a twin"
)
- parser.add_argument(
+ reference = parser.add_mutually_exclusive_group()
+ reference.add_argument(
"--expected", action="append", metavar="PATH", help="compare against this file"
)
+ reference.add_argument(
+ "--reference", metavar="SOURCE", help="compare against an explicitly selected HIR"
+ )
+ parser.add_argument(
+ "--distributed", action="store_true", help="simulate the candidate's device ranks"
+ )
parser.add_argument(
"--inputs",
metavar="random|files:A.pt,B.pt",
@@ -212,6 +219,8 @@ class CheckRequest:
expected: tuple[Any, ...] | None
device: str
expectations: dict[str, tuple[Predicate, ...]]
+ reference: Module | None = None
+ distributed: bool = False
def _refuse_orchestration(module: Module, method: str) -> None:
@@ -304,10 +313,21 @@ def _scope(resource: RuntimeResource, children: Sequence[str]) -> RuntimeResourc
def check_concrete(request: CheckRequest):
+ if request.reference is not None and (request.twin is not None or request.expected is not None):
+ raise ValueError("--reference requires a HIR candidate and cannot accompany --expected")
+ if request.distributed and request.twin is not None:
+ raise ValueError("--distributed requires a HIR candidate")
+ if request.reference is not None:
+ _compatible_reference(request.module, request.reference)
loaded = request.module.load(request.weights)
def reference_run(*args):
- return evaluate(loaded, *args)
- if request.expected is not None:
+ return evaluate(loaded, *args, distributed=request.distributed)
+ if request.reference is not None:
+ reference_loaded = request.reference.load(request.weights)
+ def selected_reference_run(*args):
+ return evaluate(reference_loaded, *args)
+ reference = selected_reference_run
+ elif request.expected is not None:
expected = request.expected[0] if len(request.expected) == 1 else request.expected
def expected_run(*_args):
return expected
@@ -322,6 +342,29 @@ def expected_run(*_args):
return check(candidate, reference, request.inputs, expect=request.expectations)
+def _compatible_reference(candidate: Module, reference: Module) -> None:
+ """Require one logical input and weight binding shared by both programs."""
+ def activations(module):
+ return tuple(
+ (param.type.shape, param.type.dtype)
+ for param in module.entry_function().params if not param.is_const
+ )
+
+ def weights(module, prefix=""):
+ result = {
+ prefix + name: (type_.shape, type_.dtype)
+ for name, type_ in module.weights.items()
+ }
+ for child in module.modules:
+ result.update(weights(child, prefix + child.name + "."))
+ return result
+
+ if activations(candidate) != activations(reference):
+ raise ValueError("reference and candidate must declare matching logical activation shapes and dtypes")
+ if weights(candidate) != weights(reference):
+ raise ValueError("reference and candidate must declare matching weight paths, shapes and dtypes")
+
+
def _walk_twin(
root: type, segments: Sequence[str]
) -> tuple[RuntimeModule, Module, Module, tuple[str, ...]]:
@@ -532,6 +575,7 @@ def _render(source: str, runs: Sequence[dict[str, Any]], warnings: Sequence[str]
if run.get("dims"):
lines += ["", " " + ", ".join(f"{k}={v}" for k, v in run["dims"].items())]
lines.append(f" reference: {run.get('reference', 'none')}")
+ lines.append(f" evaluation: {run.get('evaluation', 'local')}")
activations = run["inputs"]["activations"]
lines.append(
f" inputs: {activations['source']}; activations actual "
@@ -581,6 +625,16 @@ def run_check(arguments: argparse.Namespace) -> int:
"""Compare one selected implementation against its semantic reference."""
expect = expectations(getattr(arguments, "comparison", None))
selection = select(arguments.source)
+ distributed = getattr(arguments, "distributed", False)
+ if distributed and selection.twin is not None:
+ raise ValueError("--distributed requires a HIR candidate")
+ evaluation = "runtime" if selection.twin is not None else "distributed" if distributed else "local"
+ reference_source = getattr(arguments, "reference", None)
+ reference_selection = select(reference_source) if reference_source else None
+ if reference_selection is not None:
+ if reference_selection.twin is not None or selection.twin is not None:
+ raise ValueError("--reference selects two HIR programs, not runtime twins")
+ _compatible_reference(selection.module, reference_selection.module)
stated = parse_dims(arguments.dim) or {}
if arguments.inputs is None:
raise ValueError("no inputs stated")
@@ -625,6 +679,8 @@ def run_check(arguments: argparse.Namespace) -> int:
expected,
device,
expect,
+ reference_selection.module if reference_selection is not None else None,
+ distributed,
)
)
declared = tuple(
@@ -632,7 +688,9 @@ def run_check(arguments: argparse.Namespace) -> int:
for param in (concrete.params if concrete is not None else ())
if not param.is_const
)
- if arguments.expected:
+ if reference_source:
+ reference = reference_source
+ elif arguments.expected:
reference = ", ".join(arguments.expected)
elif selection.twin is not None:
fn_name = selection.module.entry_function().name
@@ -647,6 +705,7 @@ def run_check(arguments: argparse.Namespace) -> int:
"outputs": [_output_dict(output) for output in report.outputs],
"dims": dims,
"reference": reference,
+ "evaluation": evaluation,
"variant": selected_variant,
"inputs": {
"activations": {
diff --git a/src/tilefoundry/cli/tutorial.py b/src/tilefoundry/cli/tutorial.py
index fa9555b0..c5e2e119 100644
--- a/src/tilefoundry/cli/tutorial.py
+++ b/src/tilefoundry/cli/tutorial.py
@@ -7,7 +7,7 @@
from tilefoundry.cli import data
from tilefoundry.cli.models import render_source_directory, source_summary
-PAGES: tuple[str, ...] = ("index", "migrate", "optimize", "showcase")
+PAGES: tuple[str, ...] = ("index", "migrate", "optimize", "showcase", "distributed-check")
def page_path(page: str) -> Path:
diff --git a/tests/cli/test_cli_check.py b/tests/cli/test_cli_check.py
index ebc0e660..4e4ef49e 100644
--- a/tests/cli/test_cli_check.py
+++ b/tests/cli/test_cli_check.py
@@ -14,6 +14,7 @@
import torch
from safetensors.torch import save_file
+from tests.fixtures.distributed import projection
from tests.fixtures.placed import gqa_decode, leaf_weights
from tests.fixtures.shapes.composed_leaf_source import composed_leaf_source
from tests.models.corpus import MODELS_ROOT
@@ -707,3 +708,51 @@ def test_a_pinned_extent_on_a_root_that_reaches_a_child(tmp_path, capsys) -> Non
== 0
)
assert "PASS" in capsys.readouterr().out
+
+
+@pytest.mark.parametrize("candidate", ["TensorParallel", "ReduceScatterProjection"])
+def test_explicit_reference_checks_distributed_outputs_and_state(candidate, tmp_path, capsys):
+ reference = f"{projection.__file__}:Reference"
+ report_path = tmp_path / "distributed.json"
+ arguments = [
+ "check", f"{projection.__file__}:{candidate}", "--reference", reference,
+ "--distributed", "--inputs", "random", "--weights", "random", "--device", "cpu",
+ "--out", "output[0]", "--fn", "allclose", "--atol", "1e-5", "--rtol", "1e-5",
+ "--out", "output[1]", "--fn", "allclose", "--atol", "1e-5", "--rtol", "1e-5",
+ ]
+ assert cli.main(arguments) == 0
+ shown = capsys.readouterr().out
+ assert reference in shown and "evaluation: distributed" in shown and "PASS" in shown
+ assert cli.main([*arguments, "--json", str(report_path)]) == 0
+ assert capsys.readouterr().out == ""
+ report = json.loads(report_path.read_text())
+ assert report["passed"]
+ assert report["runs"][0]["reference"] == reference
+ assert report["runs"][0]["evaluation"] == "distributed"
+
+
+def test_explicit_reference_reports_real_disagreement(tmp_path, capsys):
+ changed = tmp_path / "changed.py"
+ changed.write_text(Path(projection.__file__).read_text().replace(
+ "return tf.neg(updated), updated", "return updated, updated", 1,
+ ))
+ assert cli.main([
+ "check", f"{projection.__file__}:TensorParallel", "--reference", f"{changed}:Reference",
+ "--distributed", "--inputs", "random", "--weights", "random", "--device", "cpu",
+ "--out", "output[0]", "--fn", "allclose", "--atol", "1e-5", "--rtol", "1e-5",
+ "--out", "output[1]", "--fn", "allclose", "--atol", "1e-5", "--rtol", "1e-5",
+ ]) == 1
+ assert "FAIL" in capsys.readouterr().out
+
+
+def test_explicit_reference_refuses_different_weight_bindings(tmp_path, capsys):
+ changed = tmp_path / "changed.py"
+ changed.write_text(Path(projection.__file__).read_text().replace(
+ 'w: ConstTensor', 'other: ConstTensor', 1,
+ ).replace('tf.matmul(x, w)', 'tf.matmul(x, other)', 1))
+ assert cli.main([
+ "check", f"{projection.__file__}:TensorParallel", "--reference", f"{changed}:Reference",
+ "--distributed", "--inputs", "random", "--weights", "random", "--device", "cpu",
+ "--out", "output[0]", "--fn", "equal", "--out", "output[1]", "--fn", "equal",
+ ]) == 1
+ assert "matching weight paths" in capsys.readouterr().err