LegalizeGraphBoundary Pass

Makes every FunctionType::Graph function legal for the host_build_graph runtime to record and replay: hoists the boundary scalars a Graph body derives, the boundary views it takes and the intermediates it allocates out to its call sites, and rejects unsupported boundaries before runtime submission.

Overview

The host_build_graph runtime records a Graph function's task topology on its first call for a Graph key and replays that recording afterwards. Construction Config values contribute to that key. Replay patches only two things: the addresses of the boundary tensors, and the values of the boundary scalars. Everything else — node count, shapes, dependency edges, block counts — is frozen into the recorded Definition.

That makes four classes of problem possible:

Problem	What the runtime does	What this pass does
A boundary scalar is derived inside the region	Explicitly extracting a value loses its parameter origin and freezes it into the recording; implicit conversion is rejected	Step A — hoists the computation to the call site, or leaves it alone when the frozen value is provably the right one
A view of a boundary tensor is taken inside the region	Freezes the first call's offset and patches only the address, so a later call reads call one's window	Step B — hoists the view to the call site, or accepts it when its window is replay-invariant
The region allocates its own intermediates	Records them correctly, on a heap it never reclaims mid-run, so the live set grows with the number of submissions	Step C — hoists the allocation to the call site as an InOut boundary tensor
The boundary itself is not cacheable	Reports an invalid Graph submission	Step D — rejects it at compile time

The first two can produce wrong answers when values lose their replay linkage. The third produces correct answers that run out of memory, or merely slower, as the layer count grows. The last is rejected at runtime; compile-time checks report statically decidable problems against the source program earlier.

Step A — derived boundary scalars

A boundary scalar carries its parameter origin in the runtime's InheritableScalar wrapper. args.scalar(k) returns that wrapper; forwarding it through add_scalar preserves the parameter index that replay refreshes. Construction values read from Config do not carry that origin; the wrapper provides no value-read accessor or implicit integer conversion. Step A moves derived computations outside the Graph body:

@pl.function(type=pl.FunctionType.Graph)
def layer(self, cur, wq, layer_idx: pl.Scalar[pl.INDEX]):
    base = layer_idx * 5120          # <- derived: no argument slot
    ...                              #    frozen at the first call's value

Step A rewrites this so the value arrives as a parameter instead:

# after the pass, conceptually:
def layer(self, cur, wq, layer_idx, base):   # base is now a real boundary scalar
    ...

# and at each call site:
self.layer(cur, wq_view(i), i, i * 5120)     # the arithmetic moved out here

A value is hoistable when its whole expression tree bottoms out in the Graph's own scalar parameters and constants — exactly the set a call site can recompute, since it already supplies those parameters. Scalar arithmetic in PyPTO is a BinaryExpr / UnaryExpr node rather than a Call, so the check recurses through those two base classes and treats everything else as a leaf.

New parameters are appended, not prepended: CoreTaskArgs requires every tensor argument to precede every scalar one.

A bare rename is deleted, not accepted

n = batch computes nothing, so there is nothing to hoist — but leaving it alone would generate int64_t n = batch;, which the runtime rejects because the wrapper has no implicit integer conversion. Explicitly extracting its value would compile, but lose the parameter origin:

const auto batch = args.scalar(0);  // preserves the parameter origin
int64_t n = from_u64<int64_t>(config.scalar_0);  // reads a construction value
g0_params_t0.add_scalar(n);         // the task no longer inherits the parameter

Recording classifies a scalar by its retained parameter origin. The integer copy no longer carries that origin, so it is recorded as static data and every later replay reuses the first call's number. Copying the wrapper itself with auto preserves its origin.

Step A therefore substitutes the name away and erases the binding, so the task reads add_scalar(batch) — the parameter wrapper. Chains collapse to their root in one pass (a = p; b = a; sends both readers to p), and an alias of a hoisted value lands on that value's new parameter. A rename that somehow survives is rejected rather than waved through, by both Step D and the verifier.

Hoistable vs replay-invariant

Not being hoistable is not the same as being illegal. Hoistability answers "can the call site recompute this?"; the runtime only needs "is this the same on every call?", which is strictly weaker. Freezing a value into the recording is a wrong answer only when the value can differ between calls.

Property	Question	Consequence
hoistable	can the call site recompute it?	Step A moves it out; it gets a real argument slot
replay-invariant	is it the same on every call?	it stays where it is, frozen — correctly

ReplayInvariantSet (utils/graph_replay_invariant.h) draws the second line. Three seeds, closed under scalar arithmetic and over names bound to an invariant value:

Seed	Why replay reproduces it
a literal	trivially
the induction variable of a constant-trip loop	recording walks the loop once and bakes each iteration's literal into that iteration's own node; constant bounds mean every later call walks the identical sequence
tensor.dim of a boundary tensor parameter, with a literal axis	graph_boundary_matches compares each boundary tensor's ndims, shapes and strides against the recorded GraphBoundarySignature and rejects a submission on any mismatch, so a boundary shape cannot change within one recording

Scalar parameters are deliberately excluded from the invariant set. The runtime patches a boundary scalar's slot on every call — which is exactly what makes one a legal task argument and an illegal frozen view offset.

This is what admits a tiled kernel at all. A slab offset i * TILE cannot be hoisted: the value does not exist at the call site. Every projection, MLP and attention loop in a decoder layer is indexed that way, so rejecting it excluded the whole shape.

DataType::TASK_ID operands are skipped outright rather than classified. A task id is never a boundary scalar — the recording captures dependency structure itself, and the runtime refuses any Graph submission carrying an explicit dependency (explicit_dep_count() != 0), so an id produced outside the region can never reach a replay. This matters because the scalar check runs over every call's arguments, so seeds[0] = seed — an array.update_element into a pl.array of TASK_ID — used to read as a task consuming a boundary scalar while the same id written straight into deps=[...] was accepted.

A scalar that reaches a task and is neither hoistable, a boundary parameter, nor replay-invariant — because it depends on a task output or a tensor read — is rejected with a message naming the variable.

Step B — derived slices of a boundary tensor

Replay patches a boundary tensor's address. A view taken inside the region is re-derived from whatever the recording froze, so it must be taken at the call site instead:

wl = pl.tensor.slice(w, [128, 128], [layer_idx * 128, 0])   # inside the region

Step B moves that slice out and passes the result in as an additional boundary tensor. Each slice site becomes its own parameter with its own fixed shape, which is what the runtime's BOUNDARY_VIEW classification requires — it matches on same-buffer plus offset, with the shape playing no part, so a view whose shape varied between calls could not be classified at all.

The hoisted statements are emitted scalars first, then tensors, because a slice's offset is typically a Step A scalar and the binding has to precede its use. The parameter order is the reverse — tensors before scalars — which is what CoreTaskArgs requires. A view of a region-local tensor stays put.

A view of a hoisted view is hoisted too. Once wl moves out it is a boundary parameter, so wr = slice(wl, ...) is in exactly the position wl was. The body is SSA in definition order, so one forward pass reaches the whole chain — a view can only name a source defined before it. Leaving wr behind is silent: graph_rebind_tensor patches the buffer address from wl but keeps the offset recorded on the first call.

Provenance crosses an in-place call. tmp = kernel(a, tmp) rebinds the buffer to a fresh SSA name, and a view of that name is still a view of a boundary tensor. Tracking only bare alias = var assignments loses the root there, and Step B then skips the view outright — no hoist and no check — so a call-varying offset stays in the region with the first call's window frozen. The rebind is followed through ExplicitReturnedParamIndices, the same return-position -> parameter map orchestration codegen aliases a call result on, so provenance and codegen cannot disagree about which buffer a result names. This applies to a boundary parameter just as much as to a Step C allocation.

The argument that map points at is resolved with CallerSuppliedArg, which splits the three regions of the Submit::args_ contract in ir/expr.h. A Submit legally omits a runtime-allocated Out tail, and its caller-supplied prefix still maps positionally, so requiring full arity would silently drop provenance for every such launch. A parameter in the runtime-allocated gap correctly yields nothing: the runtime creates that buffer, so there is no caller-side root to inherit.

A hoisted view must have a compile-time-constant shape. Replay copies a view's shapes and strides straight from the recorded template and patches only buffer_addr and start_offset, so an extent read from a boundary scalar would apply the first call's shape to a later call's buffer. An extent that is replay-invariant is accepted for the same reason it is elsewhere: it cannot differ, so freezing it changes nothing.

A view has three outcomes, not two. Replay restores a BOUNDARY_VIEW as boundary.start_offset + packed_offset, where packed_offset is the delta recorded on the first call; only the address, size, owner and version are patched. So the offset is frozen, and that is wrong exactly when it can move:

Every non-source operand is	Outcome	Why
derivable	hoisted to the call site	the caller can rebuild the view
replay-invariant only	left in place	frozen == correct, and the call site has no name for it
neither, or a mix of the two	rejected	the frozen delta would be call one's

The mixed case is the one that matters. off = layer_idx + i * TILE can neither be hoisted (i does not exist at the call site) nor frozen (layer_idx is patched per call), and it is what a tiled region produces if it indexes weights by both. Anything built on top of a view left in place — a view of it, or an alias — is held to the same rule rather than hoisted, since the call site has no name for any of them either.

A varying offset is fine, and that is not obvious. Codegen clamps a runtime view to min(declared, source.shapes[i] - offset[i]), so the actual shape is offset-dependent even when the IR extent is constant. It does not reach replay because a hoisted view is passed as its own boundary tensor rather than re-derived in the region: graph_tensor_from_boundary tries BOUNDARY_EXACT across all boundary tensors before any BOUNDARY_VIEW, and a node consuming the view matches the view itself, so graph_rebind_tensor replaces the whole GraphTensor — shapes, strides and extent_elem included. The frozen-shape behaviour belongs to BOUNDARY_VIEW, which is the in-region case this step hoists out. Requiring the offset to be provably in bounds would reject the motivating case, a per-layer layer_idx * 5120.

Step C — allocations inside the region

pl.create_tensor at the top level of the region is hoisted to the call site, where it becomes an InOut boundary tensor:

# Before
@pl.jit.graph
def layer(a: pl.Tensor, acc: pl.InOut[pl.Tensor]):
    tmp = pl.create_tensor([ROWS, COLS], pl.FP32)   # allocated per submission
    ...
    return acc

# After — the call site owns the buffer
def layer(a, acc, tmp: pl.InOut[pl.Tensor]): ...
tmp__graph_arg0 = pl.tensor.create([ROWS, COLS], pl.FP32)
acc = layer(a, acc, tmp__graph_arg0)

Recording it in place is correct — codegen lowers a create into a batched alloc_tensors and the runtime records that as a kernel-less node, the same shape submit_dummy_task records — but the buffer comes off the graph heap, which task_allocator.h never reclaims mid-run ("The whole graph must fit at once; nothing is reclaimed mid-run"). The live set then grows with the number of submissions rather than staying flat: a decoder layer holding 14 intermediates costs 14 × N over N recorded layers. Hoisting puts the buffer back on the ordinary reclaimable heap and empties the region of allocation nodes.

simpler's hand-written examples/a2a3/host_build_graph/qwen3_14b_decode scene does the same thing by hand, for the same stated reason — "This keeps the temporary live set flat in layer count and fits the default ring configuration."

The hoisted allocation parameter remains InOut: the region writes the buffer, and declared In codegen would emit add_input, the launch would never register as a writer, and a caller that hoisted the allocation out of its own loop would get no ordering between successive launches over it. Hoisting runs after direction and dependency derivation, so newly introduced allocation parameters retain conservative InOut ordering. This does not restrict existing write-only boundary parameters: their Out signature is preserved, and DeriveCallDirections selects OutputExisting for an existing first-write buffer or promotes it to InOut when ordering requires it.

A hoisted allocation is a boundary tensor, so a view of it is held to the Step B rule above rather than left alone. That is not optional bookkeeping: the GraphBoundaryLegalized verifier treats every tensor parameter as a boundary root, so a view left in place with a window that can move would make this pass produce IR its own verifier rejects.

Two allocations are deliberately left where they are:

Left in place	Why
tensor.full	Orchestration codegen has no lowering for it at the call site either, so hoisting would only move the failure. Step D rejects it outright
A create under a loop	It is a fresh buffer per iteration. Collapsing N buffers into one parameter would make iterations alias, and the cross-task edges that would have to re-serialise them were derived by AutoDeriveTaskDependencies, well upstream of here

What recording cannot reproduce either way is a shape read from a boundary scalar. The extent is copied into the node and the buffer's address derived from it, and replay never re-runs the body, so a later call with a larger extent is handed the first call's buffer — a wrong address layout rather than a fallback. Step D rejects that, before Step C runs, so every allocation Step C hoists has a compile-time-constant shape.

Prerequisite: a Graph's returns name their parameters

Step B and Step C can append writable (Out/InOut) parameters, and that is what made automating this hoist a rework rather than a wiring job. Orchestration codegen maps a call result onto one of the callee's Out/InOut params through return_lineage::ExplicitReturnedParamIndices — a pointer-identity read of the callee's ReturnStmt — and falls back to "the single Out/InOut param" only when that map yields nothing:

// GenerateSingleReturnAlias, orchestration_codegen.cpp
INTERNAL_CHECK_SPAN(returned_idx.has_value() || out_indices.size() == 1, call->span_)

A Graph body is c_1 = layer_incore_0(a, c); return c_1 once OutlineIncoreScopes has run — a rebind, not the parameter — so a Graph has always relied on that fallback, and the fallback stops existing at the second InOut parameter.

NormalizeReturnOrder supplies what the hoists need: it canonicalizes a Graph's tensor returns the way it already did for kernels and wrappers (canonicalization only — no Graph return is reordered), and IRProperty::ReturnParamsExplicit covers Graph, so the nineteen passes between there and here cannot quietly undo it. That landed separately in PR #2618, so this step depends on it rather than providing it — which is why _legalize_outlined in the unit tests runs NormalizeReturnOrder: without it the map is all-nullopt and the hoists silently do nothing.

Step D — boundary legality

Scalar-only control flow that submits no tasks may still read boundary values, for example for i in pl.range(n). Codegen places each value read during graph construction in a generated simpler::hbg::ConfigBase subclass. Its hash and equality compare the same encoded scalar bits. The runtime includes this Config in the Graph key, so a changed construction value records a distinct definition. The Config also includes tensor shapes, strides, storage sizes, alias partitions, and relative offsets within each partition. A changed boundary layout therefore selects a compatible definition; relocated buffers and uniformly shifted views retain the same key. Absolute device addresses are not part of this key. The body decodes these fields with from_u64<T>; boundary scalar handles have no value-read accessor. Task argument forwarding retains InheritableScalar handles, including floating-point parameters, so replay preserves their origin. The orchestration entry reads its scalar inputs with orch_args.scalar<T>(i), which decodes values directly in both supported runtimes.

Check	Why
The compilation targets host_build_graph	GraphTaskArgs and rt_submit_graph exist only in that runtime's orchestration API, and codegen emits them unconditionally, so a Graph built against the default tensormap_and_ringbuffer yields orchestration C++ that names undeclared symbols. Reported here, against the function the user wrote, rather than as a C++ error in generated code
At least one tensor parameter	A graph with an empty boundary has nothing to patch on replay; the runtime refuses to cache it
At most 128 tensor parameters	GRAPH_MAX_TENSOR_ARGS — the boundary is a fixed-size GraphTaskArgs
At most 64 scalar parameters	GRAPH_MAX_SCALAR_ARGS. Checked after Step A, which adds scalar parameters, so a signature that fit before hoisting can stop fitting after
No runtime-allocated boundary output	Out describes write-only access and is legal with caller-supplied storage. Call-site OutputExisting and InOut are supported; ArgDirection::Output is rejected because it requests a new allocation
Scalar parameters are In	A boundary scalar is passed by value and replayed from the call site
Returns only its own parameters	rt_submit_graph yields a valid task id only on a cache hit, so nothing can depend on a graph call's result. return c for an InOut parameter is the in-place spelling and is fine; a computed value is not
Between 1 and 1024 launched tasks	graph_execution_storage_layout refuses a node count of zero as well as one over GRAPH_MAX_NODES. A launch in a loop counts once per iteration, not once per call site, and system.task_dummy counts too — it lowers to rt_submit_dummy_task, and ExpandManualPhaseFence inserts them automatically. Allocations record nodes too, counted as an upper bound so that passing this check means the runtime will accept the Graph. Codegen collects every eligible create in a statement list — an intervening launch does not close the batch — and packs them at most kAllocTensorsArgs (16) to an alloc_tensors. Two of its three ineligibility rules cannot fire here (a shape reading a local is already rejected as non-constant; an already-declared var cannot recur under SSA), so those creates are counted exactly. The third can — an injected GM pipe buffer leaves the shared batch when its core_num reads a body-local, which only the emitter's use-resolution knows — so each of those is charged its worst case of one node. The batch size and the GM-pipe predicate are shared with the emitter in utils/alloc_batching.h rather than restated here. Counted after the hoisting steps, unlike every other check in this table: Step C removes allocation nodes, so the pre-hoist body would reject a Graph that fits and disagree with the verifier, which re-derives the same count from the rewritten IR
No allocation under a runtime loop or branch	Each records a node, so a count that varies between calls is a topology that varies
Allocation shapes are compile-time constants	Recording copies the shape into the node and derives the buffer address from it; a shape reading a boundary scalar is frozen at the first call's value
No tensor.full in the region	Orchestration codegen has no lowering for it and rejects it as a misplaced tensor op
No launch under a runtime loop, while or if	The recording fixes the topology on the first call and replays it unchanged, so a launch count or branch that can differ between calls would silently replay call one's shape
No call-varying scalar computed inline at a task argument	Step A hoists named derived values; an expression written inline at the call has no name to hoist and no boundary slot, so it would be frozen at the first call's value. A replay-invariant expression is accepted — the freeze is harmless
No Graph calls a Graph	The runtime cannot record a graph from inside one it is already recording
Every parameter supplied at the call site	A Submit may normally pass a prefix and let the runtime allocate the tail Out params; a Graph has no such tail
No explicit dependencies on the launch	The runtime rejects explicit dependencies across a Graph boundary
No dispatch predicate on the launch	A predicate on a graph launch is neither honoured nor rejected — the runtime silently zeroes it, so the region would run unconditionally

Position in the pipeline

Runs after the final Simplify and immediately before MaterializeRuntimeScopes.

That position is forced from both sides. DeriveCallDirections and AutoDeriveTaskDependencies must already have run, so argument directions and cross-task edges are known. MaterializeRuntimeScopes must not yet have run, so no scope wrapper has been placed around the statements Step A moves.

Pass properties

• Requires: SplitIncoreOrch, CallDirectionsResolved
• Produces: GraphBoundaryLegalized, CallDirectionsResolved

CallDirectionsResolved is re-declared because the pass rewrites call arguments and their direction attrs; MaterializeRuntimeScopes, which runs next, requires that property.

The GraphBoundaryLegalized verifier re-derives the topology, node count and hoisting post-conditions independently, so a later pass that reintroduces a rejected state is caught. It shares exactly one thing with this pass: ReplayInvariantSet. That is not a decision this pass made and the verifier rubber-stamps — it is a reading of the runtime's own contract, and two hand-written copies of it could disagree, leaving the verifier rejecting IR the pass had just produced.

Not yet handled

Hoisting an allocation made under a loop — it is a fresh buffer per iteration, so the call site would have to allocate an array of them rather than one (see Step C). Hoisting one whose shape reads a boundary scalar, which Step D rejects rather than rebuilding at the call site. And packing a boundary of more than 128 tensors into a scratch arena.

Neither does anything hoist a call-site allocation out of the caller's own loop: a create lands immediately before the launch, so each call still gets its own buffer. The win is that the buffer now comes off the ordinary reclaimable heap instead of the graph heap, which is independent of where the caller puts it.

See also

• Pass Manager — full pipeline order
• MaterializeRuntimeScopes — runs immediately after

Host control tensors

With host_build_graph, orchestration-level tensor.read and tensor.write operate on ready CPU tensor arguments. Generated argument metadata selects HOST/NONE transfers for those arguments; ordinary device tensors retain their normal upload behavior. This metadata is preserved by saved artifacts and worker reuse, including distributed chip dispatch. Scheduler predicate operands remain device tensors and use the normal device transfer path. Control tensors must stay alive until the call completes.

Distributed dispatch also accepts an existing simpler.Tensor backed by HOST storage for a host control argument. The address-free descriptor is reused with TensorTransfer.NONE. Keep the owning Buffer and any fork-inherited storage alive until completion. A wire Tensor snapshots its view metadata, so construct another view when the shape, strides, or storage changes.

A control tensor (including a view or loop-carried alias) cannot also be passed to a device kernel through the same argument. Pass host control data and device data in separate parameters. Host writes do not upload to device storage. Runtime allocations and values produced by device tasks in the current graph cannot be read during host graph construction. Passing a DeviceTensor for a host control parameter raises an error before dispatch.