Skip to content

MaterializeCommDomainScopes Pass

Overview

MaterializeCommDomainScopes walks each host-orchestration function and assembles the host-side metadata that the distributed runtime needs in order to size and populate per-rank communication windows. It is the structural analogue of InitMemRef: it traces an allocation through to its consumption points, constructs a back-reference object, and threads it onto the IR types so downstream codegen has O(1) access.

Aspect MemRef side WindowBuffer side
Allocation op tile.alloc(memory_space, size_in_bytes) pld.tensor.alloc_window_buffer(size_in_bytes)
Assignment LHS at parse time Var(PtrType) Var(PtrType) (same singleton)
Wrapper Var subclass MemRef WindowBuffer
Wrapper's SSA-edge type MemRefType (singleton) WindowBufferType (singleton)
Built by InitMemRef MaterializeCommDomainScopes (this pass)
Threaded back onto TensorType.memref_ DistributedTensorType.window_buffer_
IR-level registry Program.functions_ (alloc stmts) CommDomainScopeStmt wrapping each host_orch body (one per inferred comm domain)

Position in the pipeline

... -> ExpandManualPhaseFence -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> MaterializeDistTensorCtx -> Simplify (final)

The pass runs near the end of the default pipeline, immediately before LowerHostTensorCollectives and the final Simplify. None of the intervening passes between InlineFunctions and here touches the host_orch alloc/window/dispatch chain: host_orch is never tile-lowered, and L2 (chip-level) orchestrations are never inlined into L3, so the alloc / view / dispatch sites this pass needs are still discoverable. Running late keeps the producing IR fully canonicalised before the descriptor analysis kicks in, and any constant folding that the trailing Simplify does on the collected sizes is applied uniformly.

Algorithm

For every host-orchestration function (Function::level_ == Level::HOST and Function::role_ == Role::Orchestrator, regardless of func_type_):

  1. Collect allocations. Find every AssignStmt whose RHS is a pld.tensor.alloc_window_buffer(size, *, name) Call. Record (ptr_var, size_expr, name, span, call).

  2. Collect views. For every AssignStmt whose RHS is a pld.tensor.window(ptr_var, [shape], *, dtype) Call referencing a recorded ptr_var, record the binding view_var → alloc.

  3. Scan dispatches. Walk the body with a stack of enclosing ForStmts. For every Call whose op_ is a GlobalVar resolving to a chip-level orchestration, read attrs["device"] and infer a device descriptor from the device expression in the current loop context:

device= shape Descriptor
ConstInt(N) subset = {N}
IterArg of for r in pl.range(pld.system.world_size()) kAll
IterArg of for r in pl.range(ConstInt(N)) subset = {0, …, N − 1}
other pypto::ValueError

Every positional dispatch arg that is a recorded view Var contributes that descriptor to its underlying allocation.

  1. Merge descriptors. Per allocation, fold every recorded descriptor: any kAll ⇒ kAll; otherwise union the subsets.

  2. Materialise WindowBuffers. For each allocation construct WindowBuffer(base = ptr_var, size = size_expr, load_from_host = false, store_to_host = false). The Var::name_hint_ is inherited from ptr_var->name_hint_. (Host-staging flags are placeholders for N4+.)

  3. Rewrite view types (host_orch only). For every view binding, mint a fresh Var of the same name_hint_ whose type is DistributedTensorType(shape, dtype, memref, tensor_view, wb) and run Substitute to swap every reference to the old view Var with the fresh one. Two pld.tensor.window views over the same allocation share the same shared_ptr<const WindowBuffer>. Chip-orch / InCore parameter types are not touched.

  4. Cluster into comm domains. Walk the allocation list in source order; append a slot to the first existing comm domain whose merged device descriptor matches, or open a new one.

  5. Wrap the body in scope statements. Build a chain of nested CommDomainScopeStmts — one per comm domain, outer = first declared, inner = last — and substitute it for the host_orch function body. Each scope carries its devices list and its slot vector; name_hint_ is set to "comm_d<n>" so codegen emits the matching __comm_d<n> handle variable verbatim.

Sanity checks

The pass raises pypto::ValueError (carrying the alloc's span) if:

  • An allocation has no pld.tensor.window materialisation (dead alloc).
  • An allocation has at least one view but no chip-orch dispatch consumes it.
  • The device= expression on a dispatch is something other than ConstInt or a recognised pl.range induction var.
  • Two allocations within the same comm domain share a name_hint_ (the parser already enforces global uniqueness; the pass re-asserts).

Output invariants

After the pass:

  • Every host_orch function whose body contains at least one alloc has its body wrapped in a chain of nested CommDomainScopeStmts (one per inferred comm domain, outer = first declared, inner = last). Allocation-free host_orchs are left unchanged.
  • Every pld.tensor.window result Var's type is a DistributedTensorType whose window_buffer_ field points to the corresponding WindowBuffer.
  • Every host-level pld.tensor.allreduce call has two positional arguments after SynthesizeAllReduceSignals runs. For an omitted user signal, the second argument is a synthesized Var produced by a preceding pld.tensor.window assignment.
  • pld.tensor.window views over the same allocation share the same shared_ptr<const WindowBuffer> — pointer-equality is a load-bearing invariant for downstream codegen.
  • Chip-orchestration and InCore parameter types remain nullopt on window_buffer_. N7 codegen reads the back-reference at the host_orch dispatch site and threads the matching CommContext pointer explicitly.

Pass properties

Field Value
required {}
produced {IRProperty::CommDomainScopesMaterialized}
invalidated {}

Reference