feat(hir): settle where cross-device communication semantics live

[zhen8838]

@UranusSeven I read through #178 — the RFC and the rank-local evaluator are good work, and the evaluator in particular is what issue #138 has been asking for. A few things I'd like to discuss.

1. I'd like the value types unified

A few earlier PRs already laid some multi-card groundwork that the plan's Current state doesn't reference, so listing them here:

• refactor(runtime): close the gaps PR #161's review and discussion turned up #162 — a mesh is stated as one CuTe-style layout, and one mesh may name several levels
• feat(runtime): load each program its own shard of a checkpoint #165 — each program loads its own shard of a checkpoint; Placement tells it which program it is; ShardTensor.to_local produces its local slice
• feat(shard): express swizzled composed layouts as CuTe does #174 — swizzled composed layouts

ShardTensor (runtime/tensor.py:27) answers exactly "which slice is mine", and to_local runs the same address algebra the device side runs. DistributedValue plus _slices / _attrs / _local_type is a second, independent slicing. When two independent slicings disagree, that disagreement is precisely what check exists to catch — which makes it an awkward thing to build check on top of.

I'd like these unified. If the runtime → evaluator dependency direction is wrong, ShardTensor can move somewhere both can reach.

2. I'd suggest dropping --distributed

The program's mesh already states whether it is distributed. A separate flag lets the two disagree — a multi-card program without the flag silently takes the single-participant path, which is the unsoundness issue #138 describes.

Beyond that, per-unit execution isn't specific to multiple cards. After #175 a loop start may contain a MeshCoord, so two CTAs can have different iteration domains; the same holds at thread level. The criterion is independent of topology level, and #138 already states it:

unit_varying(e) = e is Local on a Split axis | any operand is unit_varying

Local is constructed in exactly one place (parser/pattern_nodes.py:2903), and #138 measured the reach: 55 of 814 Call nodes across the fixture corpus, all inside the two functions that read a coordinate. Letting the evaluator decide for itself covers gpu, cta and thread alike.

There's a second layer I'd like to discuss along with it: can the evaluator decide how many cards to actually use? Today it replays coordinates(mesh) sequentially on one device. That's fine for small programs, but once a model is large, N full tensors on one card neither fit nor finish in reasonable time. When the machine has several cards, the evaluator should be able to place those ranks on real devices rather than only simulate them. Same direction as above — whether per-unit execution is needed, and where those units go, are both things the evaluator can determine rather than command-line switches.

--reference is different: a second HIR cannot be inferred from the program, so that argument earns its place.

3. New ops are fine — I'd like search and evaluation to face only reshard

This is the one I most want to discuss, so the conclusion first: collective ops can exist. What I'd like is for the search and evaluation stages to need only reshard. reshard is the logical form; a distributed search materializes it into concrete collectives later.

A few reasons that layering appeals to me.

A smaller and more comparable candidate space. The author writes a layout change; which collective implements it is a materialization decision. All four collectives are already one ShardLayout transition in the existing types:

Split(k)  -> Broadcast     all-gather
Partial   -> Broadcast     all-reduce
Partial   -> Split(k)      reduce-scatter
Split(i)  -> Split(j)      all-to-all

The last one already runs in tests/fixtures/placed/tp_all_to_all.py, whose entire body is tf.reshard(x, (R, C @ g.g), "gmem").

analyze works today on that form. The PR notes that "all existing analyze selectors explicitly reject uncosted collectives", so a candidate written with collectives cannot be analyzed — and the RFC's check → analyze → rank loop breaks at the second step. Meanwhile tp_all_to_all.py's docstring says the program "exists so an analysis can say what it would cost, which is what tests/analysis/test_analysis_families.py asks of it".

Ordering doesn't have to be answered at this level. RFC §6.2 says collective ordering "must be settled before collective implementation", and the plan's D3 currently relies on operand-order traversal. If collectives are a materialization product, ordering belongs to materialization and HIR doesn't have to settle it first.

One concrete note: _reshard currently refuses cross-device movement (Reshard requires cross-device data; use an explicit collective), so tp_all_to_all.py — an existing fixture — cannot run under --distributed.

I agree with RFC §6.2 that "a layout annotation alone cannot cause an unreported transfer". The resolution I'd reach for is that a reshard crossing the gpu axis must appear in traffic and cost (#168 and #177 are building that), rather than a second way to spell it.

Also, the two constraints in collective.py — a single-level gpu mesh and contiguous coordinate-order strides — exclude the multi-level meshes from #162 and the swizzled layouts from #174. I'd like to know whether those are requirements of the semantics or of the simulation.

RFC §12 lists this as its first open question and §10 Phase 1 says to resolve those before implementing affected surfaces, so this seemed like the right moment to raise it.

Questions

• How do you see the layering — search and evaluation stopping at reshard, with collectives chosen at materialization?
• Can the evaluator use several real cards when the machine has them, instead of replaying N ranks on one?
• Do the four ShardLayout transitions cover the primitive set §12 asks about? If not, which case is missing?
• Are the two mesh constraints in collective.py required by the semantics or by the simulation?
• Could ShardTensor serve as the evaluator's value type, and where should it live if so?