InsertCommFence Pass¶

Overview¶

InsertCommFence implements the data-before-signal memory-consistency contract the latest PTOAS defines in its pto-memory-consistency pass and pushes onto the compiler. "Pushes onto the compiler" is literal: ptoas requires the markers but does not check for them, so satisfying the contract is entirely this pass's job (see ptoas does not gate this). The contract is two-sided:

Publish side. A pto.comm.tnotify requires an explicit pto.fence.barrier_all #pto.fence_scope<gm> after the matching pto.cmo.cacheinvalid release marker and before the signal — a cross-rank write must be visible to the peer before the notify that signals it.
Consume side. A cacheable GM load after pto.comm.twait (or a successful pto.comm.ttest) requires an explicit pto.cmo.cacheinvalid all #pto.address_space<gm> first, so the reader sees the peer's fresh write.

That ttest half of the consume side is a statement about the hardware contract, not about this pass: PyPTO has no ttest operator, so pld.system.wait is the only consume-side point the pass can recognise. If a ttest op is ever added, StmtEffect must classify it as Effect::kWait too.

Both cache markers are the same system.cacheinvalid op, in two forms distinguished by arity:

Form	IR	Lowers to
Region	`system.cacheinvalid(tensor, shapes, offsets)`	`pto.cmo.cacheinvalid … single_cache_line`
Whole-GM	`system.cacheinvalid()` (no args)	`pto.cmo.cacheinvalid all #pto.address_space<gm>`

system.fence lowers to pto.fence.barrier_all #pto.fence_scope<gm> — a GM barrier with DDR observability, stronger than a bare pto.barrier <PIPE_ALL>.

What the pass inserts¶

Verified empirically on ptoas 0.50, the contract reduces to two purely-local rules — the notify itself needs no marker. A single structural traversal (InsertCommMarkers), with no control-flow state, inserts:

after each local publishing write — a tile.store or tensor.write into a window-bound destination, or a get into a window-bound local destination: a whole-tensor region system.cacheinvalid of the written target, immediately followed by a system.fence;
after each remote publishing write — remote_store / put: only a system.fence (see the note below on where its cacheinvalid comes from);
after each opaque publishing write — a Submit, or a call to an unregistered user function whose body is not analysed here (no single addressable region): a conservative whole-GM system.cacheinvalid() + system.fence;
after each wait — a whole-GM system.cacheinvalid (the consume-side invalidate before the next cacheable read). Batched: a run of consecutive waits, or a pure wait-loop (a for/while whose body — through if/seq nesting — contains only waits and at least one, e.g. the mesh composite's for src: if src != me: wait(...)), performs no memory access between the waits, so ONE whole-GM cacheinvalid is emitted after the run/loop instead of one per wait — (P-1) whole-cache flushes become 1 per barrier generation. The traversal stays structural (no dataflow), using a local pure-wait-loop check; loops that also load (e.g. ring's per-step wait; load) are not pure and keep the per-wait invalidate. A wait-free loop is not pure either: an empty body is vacuously "wait-only", so the classification requires ≥ 1 wait — otherwise the pass would append a whole-GM cacheinvalid after a loop that emitted nothing before (a regression in a pass whose point is removing flushes);
notify — nothing.

A region system.cacheinvalid(target) addresses target's local base, which is correct for a local-window store. The remote writes remote_store / put write to a peer-offset GM address (local_ptr + delems(peer)) that the local target view does not address — a local-target cacheinvalid would clean the wrong line. The peer offset is only known during codegen (EmitCommRemoteView) and is not yet expressible in the IR, so the peer-region pto.cmo.cacheinvalid <peer_view> single_cache_line is emitted by the op's codegen as a workaround. The GM release fence, however, is always an explicit system.fence op inserted by this pass — the codegen must not embed it, keeping the release ordering IR-visible and uniform. (Giving the peer-region cacheinvalid a first-class IR representation, so the pass can own the whole marker, is a follow-up.)

store(win_a); store(win_b); notify        (local window stores)
  -> store(win_a); cacheinvalid(win_a); fence; store(win_b); cacheinvalid(win_b); fence; notify

for c: store; for p: notify                 (write and notify in separate loops)
  -> for c: (store; cacheinvalid; fence);
     for p: notify

wait; read
  -> wait; cacheinvalid(); read

Why the notify needs no marker, and the fence sits at the write: ptoas associates the required release fence with a publishing write's cacheinvalid, not with the notify. So a tnotify that releases data written earlier — even in a different loop — is already satisfied by that write's cacheinvalid; fence; the fence does not need to sit next to the notify. A pure barrier notify (no data at all) requires nothing. (This was verified by removing the notify-side markers from the ring-allreduce .pto and confirming ptoas 0.50 still accepts it; removing the wait-side cacheinvalid all, by contrast, was rejected on 0.50.)

ptoas does not gate this — the markers are the compiler's obligation¶

ptoas 0.52 accepts a module with the markers missing. Re-running the check above on 0.52, against the same ring_step kernel from tests/st/distributed/collectives/test_l3_allreduce_ring.py:

Variant	ptoas 0.52
unmodified	accepted
wait-side `cacheinvalid all` removed (rejected on 0.50)	accepted
every `cacheinvalid` and `system.fence` removed	accepted

Nothing is diagnosed; the corresponding instructions are simply not emitted:

IR marker	Instruction lost when absent
region `cacheinvalid`	`PTOAS__DCCI_SINGLE_CACHE_LINE(<GlobalTensor>)`
whole-GM `cacheinvalid()`	`dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)`
`system.fence`	`pipe_barrier(PIPE_MTE2/MTE3/FIX); dsb(DSB_DDR)`

So there is no compile-time safety net. A missing marker is not a build error but a data race: the reader may observe a stale cache line, dependent on timing, cache state, transfer size and rank count. Do not read a green test run as evidence that a marker is unnecessary — the pass's own history shows the failure mode is silent. While the pin was on ptoas 0.50 the publish-side region cacheinvalid emitted no call at all (hw-native-sys/PTOAS#995), so it never reached the device, and the distributed suite stayed green throughout.

The region cacheinvalid currently covers the whole target tensor (region [0, …] offsets over the tensor's full shape), reusing the tensor type's dim exprs. Narrowing to the precise written sub-region — the write's own (shapes, offsets) are right there at the write site — is a planned follow-up.

Markers land where the write / wait / notify is (always in scope)¶

Placing the region cacheinvalid immediately after its write means the target Var is trivially in scope (the write just used it) — whether it is a window parameter, an alias (dv = pl.tensor.view(win); remote_store(dv)), a loop-carried iter_arg, or a value defined inside a branch. There is no cross-scope tracking and nothing is ever silently dropped: every marker lands next to the op that needs it, at every nesting level. A bare single-statement branch/loop body is wrapped in place (body -> { body; markers }); after the first run the body is a SeqStmts, so the pass stays idempotent.

Position in the pipeline¶

... -> ClassifyIterArgCarry -> InsertCommFence -> MaterializeValidShapeSymbols   (last)

It runs after every statement-reordering pass in the Default pipeline (SkewCrossCorePipeline, LowerPipelineLoops, CanonicalizeIOOrder, ...). The inserted ops have no operands and no dependency edges, so an earlier insertion could be moved away from its notify/wait; running here keeps them adjacent through codegen. Only MaterializeValidShapeSymbols follows, and it rewrites device-kernel signatures rather than reordering statements. The passes before it only touch orchestration bodies (Orchestration and Graph), so the InCore IR this pass sees is exactly what codegen lowers.

Which writes the pass marks¶

Case	cacheinvalid	fence
`tile.store` into a window-bound `DistributedTensorType` (a peer can `remote_load` it)	pass (local region, dst arg 2)	pass
`tensor.write` into a window-bound `DistributedTensorType` (scalar write, kept unconverted by `ConvertTensorToTileOps`)	pass (local region, dst arg 0)	pass
`pld.tile.get` / `pld.tensor.get` into a window-bound `DistributedTensorType` destination	pass (local region, dst arg 0)	pass
`pld.tile.remote_store` / `pld.tile.put` / `pld.tensor.put` (peer-offset write)	codegen (peer region — IR workaround)	pass
`Submit`, or a call to an unregistered user function (no single addressable region)	pass (whole-GM, no args)	pass

A remote_load (result is a tile, no GM write) and a tile.store / tensor.write / get whose destination is a plain Tensor rather than a window-bound DistributedTensor are not publishing writes — no marker at all.

Algorithm — one structural traversal, with consume-side batching¶

The pass carries one piece of control-flow state — a flag that suppresses the per-wait invalidate while visiting the body of a pure wait-loop (a for/while whose body contains only pld.system.wait through seq/if nesting — at least one — with memory-inert control expressions; a wait-free loop is not a pure wait-loop):

at each local publishing write, append region cacheinvalid; fence;
at each remote publishing write (remote_store / put), append fence only (the peer-region cacheinvalid is codegen's — see below);
at each wait, append cacheinvalid() — batched: a pure wait-loop shares ONE whole-GM cacheinvalid() after the loop, and a run of consecutive waits shares ONE after the run. No memory access occurs between the waits, so an invalidate after the last is equivalent to one after every wait — (P-1) whole-cache flushes per barrier generation become 1 in the mesh composite.
notify is left untouched.

if/for/while bodies are visited normally; the only special handling is wrapping a bare single-statement body (a write/wait that is the sole body of an if/for without an enclosing SeqStmts) so its marker still lands — a bare pure wait-loop body gets the same suppression and one cacheinvalid() after the loop. Because the rules are otherwise local and append-only, control flow is irrelevant: a write inside one loop is correctly marked whether or not its notify lives in another. A control expression that reads GM (tensor.read lowers to a cached load) disqualifies a loop from being pure, so the invalidate is never deferred past a read that could observe stale peer data.

store(win); notify                   -> store(win); cacheinvalid(win); fence; notify
store(win); for: notify               -> store(win); cacheinvalid(win); fence; for: notify
for: { notify; store(win) }           -> for: { notify; store(win); cacheinvalid(win); fence }

An existing region cacheinvalid immediately followed by a fence after a write, and an existing whole-GM cacheinvalid immediately after a wait, are recognized and not duplicated, so the pass is idempotent.

Codegen interaction¶

remote_store and put codegen each emit their peer-region pto.cmo.cacheinvalid <peer_view> right after the store (the peer offset is only known there and is not yet expressible in the IR — a workaround; the pass supplies the paired GM fence as a system.fence op, so it does not emit the fence here). put additionally keeps a tail pto.barrier <PIPE_ALL> between the TPUT and its cacheinvalid: TPUT is a DMA, and the GM fence orders memory but does not drain the MTE pipe that issued the DMA, so without this barrier the following notify can fire before the (possibly atomic) TPUT has landed at the peer — test_l3_put atomic-add / subregion flake on device without it, and device-testing shows an MTE3-scoped barrier is not enough (only PIPE_ALL is stable). A PTOAS#872 workaround; remove once PTOAS drains the tput itself. The TPUT/TGET pre barriers and the TGET tail barrier are likewise kept.

Follow-up: the peer-region cacheinvalid for remote writes is the one marker still emitted in codegen rather than as an IR op, because the peer address (local_ptr + delems(peer)) has no IR representation yet. A first-class IR form would let InsertCommFence own the complete remote-write marker (cacheinvalid + fence), matching the local-write path.

Consumers¶

None downstream in the pipeline. PTO codegen lowers the inserted system.cacheinvalid (region or whole-GM by arity) and system.fence via their existing op handlers; no other pass reasons about them.