Skip to content

Scopes and Placement

Where work runs: marking a region as device work, grouping co-scheduled cores, and spreading a kernel across blocks.

Prerequisites: Functions and Programs and Programming Model § execution model.

Concept

Placement answers one question: which piece of hardware runs this code.

There are four constructs, all written with with (or for), and they compose:

Construct Places work on
pl.at A core group — marks a region as device work
pl.cluster One physical cluster — co-schedules a Cube and a Vector kernel
pl.spmd n blocks — the same kernel, once per block
pl.split_aiv Two AIV lanes — splits one region across both

The alternative to pl.at is writing a separate @pl.jit.incore function and calling it. They produce the same thing: pl.at is outlined into exactly such a function during compilation. Use the scope when the region is short and belongs where it is written; use a separate function when it deserves a name or is called from more than one place.

Placement is not the same as ordering — what must finish before a task starts. The runtime derives that from the parameter directions in Types and the buffers each task touches. That machinery, and the interfaces for steering it by hand, are Tasks and Ordering; this page is only about where code lands.

Quickstart: mark a region as device work

import pypto.language as pl
import torch
from pypto.runtime import RunConfig

CFG = RunConfig(platform="__PLATFORM__")
torch.manual_seed(0)
X = torch.randn(256, 128, dtype=torch.float32)
Y = torch.randn(256, 128, dtype=torch.float32)
@pl.jit
def scale(
    x: pl.Tensor[[256, 128], pl.FP32],
    out: pl.Out[pl.Tensor[[256, 128], pl.FP32]],
):
    with pl.at(level=pl.Level.CORE_GROUP):
        out[:] = pl.mul(x, 2.0)
    return out


out = torch.zeros(256, 128, dtype=torch.float32)
scale(X, out, config=CFG)
torch.testing.assert_close(out, X * 2.0, rtol=1e-4, atol=1e-4)
Element What it does
@pl.jit The entry point — control plane, which cannot hold operators itself
with pl.at(level=pl.Level.CORE_GROUP) Marks the region as device work, giving the operators somewhere legal to live
pl.mul(x, 2.0) Runs on the core group

Without the pl.at, this kernel fails to compile with Misplaced tensor op ... should be inside InCore block — see Functions and Programs.

Mechanics

pl.at

level= picks the hierarchy level. pl.Level.CORE_GROUP is the one that produces an InCore scope, and the region becomes its own kernel function during compilation.

Two optional keywords shape that outlined kernel:

Keyword Meaning
optimizations=[pl.split(mode)] Cross-core split mode for the outlined kernel
optimizations=[pl.cross_core_slot(slot_num=N)] Ring depth of the automatic cross-core pipeline
name_hint="..." Name for the outlined function

Entries in optimizations= must be written inline at the call site — the parser reads the AST, so a list built up in a variable is not accepted. pl.split and pl.cross_core_slot are orthogonal and combine freely: one splits the work, the other sizes the channel.

with pl.at(level=pl.Level.CORE_GROUP,
           optimizations=[pl.split(pl.SplitMode.UP_DOWN),
                          pl.cross_core_slot(slot_num=4)]):
    ...

Omitting cross_core_slot keeps the default ring depth of 2 slots per active direction — enough to double-buffer the handoff while leaving on-chip room for the tiles themselves. Raise it when the producing core should be able to run further ahead.

SPMD

pl.spmd(n) runs the same kernel on n blocks. Two forms, differing in whether the body reads the block index:

# Dispatch form — the body launches a kernel defined elsewhere. `self.kernel`
# means this form needs @pl.program; from @pl.jit, use the loop form below.
with pl.spmd(4):
    out = self.kernel(a, b, out)

The loop form is the one a @pl.jit entry can write directly:

@pl.jit
def spmd_add(
    a: pl.Tensor[[256, 128], pl.FP32],
    b: pl.Tensor[[256, 128], pl.FP32],
    out: pl.Out[pl.Tensor[[256, 128], pl.FP32]],
):
    for i in pl.spmd(2):                    # `i` binds the block index
        off = i * 128
        out = pl.store(
            pl.add(pl.load(a, [off, 0], [128, 128]), pl.load(b, [off, 0], [128, 128])),
            [off, 0],
            out,
        )
    return out


out = torch.zeros(256, 128, dtype=torch.float32)
spmd_add(X, Y, out, config=CFG)
torch.testing.assert_close(out, X + Y, rtol=1e-4, atol=1e-4)

A with pl.spmd(n): body that neither reads the block index nor dispatches a kernel is rejected — every block would be doing identical work.

When a hard pl.system.syncall is involved, size the launch from the device rather than from a literal: pass pl.system.available_cluster_count() (mixed or cube-only kernels) or pl.system.available_aiv_count() (vector-only), written inline at the call site.

Clusters and AIV lanes

with pl.cluster(): groups AIC and AIV kernels so they are co-scheduled on the same physical cluster, producing a Group function.

for aiv_id in pl.split_aiv(2, mode=...): splits one region across the two AIV lanes. It belongs to mixed-kernel programming — AIC and AIV cooperating inside one function — which the tutorials chapter covers end to end.

mode= picks how the two lanes divide the work:

mode= Each lane gets
pl.SplitMode.UP_DOWN / LEFT_RIGHT Half of every tile (rows / cols) — data-parallel
pl.SplitMode.NONE The full body; you dispatch disjoint work via aiv_id — task-parallel

Opening one region changes the rules for the whole function. The regions then own every placement decision for vector work, so vector compute has to live inside a region:

with pl.at(level=pl.Level.CORE_GROUP):
    for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
        ...                                    # phase 1 — per-lane work via aiv_id
    pl.system.syncall(core_type=pl.KernelType.MIX)  # barrier: outside, runs on both
    mm = pl.matmul(q, k)                       # cube work: outside, runs on AIC
    for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
        out = pl.add(pl.aiv_shard(mm), bias)   # phase 2 — full-width vector work

One region per vector phase. mode=NONE is the wrapper for a phase you do not want halved: both lanes run the full body, which is what un-regioned vector code did anyway, so wrapping it changes the text and not the execution. Cube ops and barriers stay outside.

mm is cube-produced and read on the vector lane, so it crosses the AIC/AIV boundary — pl.aiv_shard is what says so. The next section explains why that is required.

A function with no pl.split_aiv at all is unaffected — write it exactly as before.

Every crossing in one function must agree on split-vs-no-split. All the pl.aiv_shard / pl.aic_gather calls in a function ride a single cross-core pipe, and the hardware fixes that pipe as either split or un-split for its whole lifetime. So a mode=NONE region that crosses the boundary cannot sit beside an UP_DOWN or LEFT_RIGHT region that also crosses it:

for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    a = pl.exp(pl.aiv_shard(mm0))         # crossing, no split
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
    b = pl.exp(pl.aiv_shard(mm1))         # crossing, split      -> rejected

Two different split axes are fine, because the axis is chosen per transfer — only split-vs-no-split belongs to the pipe:

for r in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
    a = pl.exp(pl.aiv_shard(mm0))
for c in pl.split_aiv(2, mode=pl.SplitMode.LEFT_RIGHT):
    b = pl.exp(pl.aiv_shard(mm1))         # accepted

A region that carries no crossing is free to use any mode — the mode=NONE region that only pins a pld.system.notify to the vector lane never touches the pipe. When two phases genuinely need different transports, put them in separate pl.at(level=pl.Level.CORE_GROUP) scopes: each becomes its own function, and so gets its own pipe.

Name every tile that crosses a region edge

Once a function opens a region, a tile crossing a region edge must say so. The boundary between the two cores is yours to place in manual mode, so the compiler stops choosing it for you:

Direction Where it is written Op
Cube value read on the vector lane (C->V) at the top of the region pl.aiv_shard(x)
Vector value read on the cube lane (V->C) inside the region, before the read pl.aic_gather(x)
mm = pl.matmul(q, k)                        # cube, outside every region
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    v = pl.exp(pl.aiv_shard(mm))            # C->V: named
    kv = pl.aic_gather(v)                   # V->C: named
out = pl.matmul(kv, w)                      # cube again, outside

Drop either call and the crossing still works — the compiler emits the same transfer either way — but it is then a boundary nobody chose, in a program whose whole point is that you chose it. So it is rejected instead; the diagnostic names the value and the op that reads it.

In a mode=NONE region both ops cross without splitting. There is no split axis to halve or re-join, so the shape passes through unchanged — pl.aiv_shard of a [128, 128] tile is a [128, 128] tile. Only in UP_DOWN / LEFT_RIGHT do they also halve (shard) and re-join (gather).

Only gather a value both AIV lanes agree on. The hardware requires both sub-lanes to take part in a no-split handshake, and they share one destination slot with no per-lane offset. Nothing arbitrates between them: both lanes push, so if they hold different values the cube receives an unspecified one of the two. Not lane 0's — unspecified.

There is no way to select a lane here. Guarding the production of the value does not help: lane 1 still reaches the push and still sends whatever its tile holds. So a pl.aic_gather out of a mode=NONE region is well-defined only when the value is lane-uniform — computed identically on both lanes, or made identical before the gather. If the lanes must contribute different data to the cube, this construct cannot express it; route it through GM and order it yourself, or use a data-parallel (UP_DOWN / LEFT_RIGHT) region, where each lane owns a declared half and the gather re-joins them.

The compiler does not check any of this.

GM traffic is not covered by any of this. These rules are about tile values crossing a region edge. A GM tensor belongs to no lane, so no boundary op can express a crossing through one — pld.tensor.put takes a GM tensor by signature. AIC and AIV run asynchronously. ExpandMixedKernel handles one narrow C->V case automatically: a unique cube tile.store producer whose same-origin vector tile.load is in the same body or a nested body. Every other GM handoff — including V->C, communication ops, and sibling bodies — stays yours. syncall alone only aligns arrival: publish the producer's cache lines and issue a GM fence before the barrier, then invalidate the consumer's cache before it reads. For a buffer that may span multiple cache lines, use the conservative whole-GM pl.system.cacheinvalid() form; the tensor-region overload currently covers only the cache line containing the view's base address.

Put cross-rank comm ops in a region

A region also decides placement for ops that have no lane of their own. pld.system.notify is core-agnostic — the hardware runs TNOTIFY on either core — so in a kernel that mixes cube and vector work the compiler emits it on both the AIC and the AIV lane. Wrap the comm phase in a region and it is pinned to the vector lane instead:

for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    pld.tensor.put(dst=win, peer=peer, src=out,
                   dst_offsets=[0, 0], src_offsets=[0, 0], shape=[16, 256])
    pld.system.notify(target=sig, peer=peer, offsets=[0, 0], value=1,
                      op=pld.NotifyOp.AtomicAdd)

The cube copy is the dangerous one: the AIC lane can reach the notify before the AIV lane's put has landed the data, publishing a signal for bytes that are not there yet. A region removes it.

Shard once-only side effects across the two AIV lanes

A mode=NONE region body runs on BOTH AIV sub-lanes. That is the whole point of the mode — the region is not "one lane", it is two lanes running the same code, and you dispatch the disjoint work with the loop's aiv_id. The snippet above is therefore still incomplete: it fires one notify per lane, i.e. two notifies for the same peer.

An op whose side effect must happen once per logical occurrence — a pld.system.notify above all — must be either sharded by aiv_id or guarded to one lane:

# sharded: each lane takes a different set of peers
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    for owner in pl.range(aiv_id, NUM_PEERS, 2):
        pld.system.notify(target=sig, peer=owner, offsets=[0, 0], value=1,
                          op=pld.NotifyOp.AtomicAdd)

# guarded: lane 0 does it, lane 1 skips
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.NONE):
    if aiv_id == 0:
        pld.system.notify(target=sig, peer=peer, offsets=[0, 0], value=1,
                          op=pld.NotifyOp.AtomicAdd)

The guarded form carries an ordering obligation the sharded one does not. The two AIV lanes run asynchronously — nothing orders lane 0 against lane 1. So a lane-0-guarded notify may publish before lane 1's writes have landed, and the peer then reads data that is still in flight. It is safe only when the data the signal releases was written by lane 0 itself. If lane 1 contributes any of it, order the two lanes explicitly before the notify, or prefer the sharded form — there each lane releases only what it wrote, so the question does not arise.

What goes wrong without it. NotifyOp.AtomicAdd accumulates into the peer's slot. Two lanes notifying the same peer make that rank's counter read 2 when one rank has arrived. A pld.system.wait waiting for two arrivals is released by one — so that rank runs ahead and reads a buffer whose data has not landed. The symptom is wrong numbers on one rank, intermittently, at multi-rank runtime; it does not reproduce on one rank and it does not look like a synchronisation bug.

What the compiler does and does not do here

Behaviour What it means
Does Keeps a region's comm ops off the cube lane. Outside a region they are duplicated onto the AIC lane as well.
Does not Check the lane-sharding. A notify that both AIV lanes run against the same peer compiles cleanly and is not diagnosed.

The compiler cannot diagnose it: the correct form and the wrong form produce the same single statement in the AIV function, differing only in whether aiv_id reached the call's arguments. Getting this right is the author's job.

Cross-lane ordering is yours too

AIC and AIV run asynchronously. A boundary op orders the one value it carries — that is what the transfer is — but nothing orders a cube-lane write against a vector-lane read of the same GM buffer. Publish the producer's cache lines and issue a GM fence, place a cross-core barrier between the phases, then invalidate the consumer's cache before it reads. The barrier by itself synchronizes arrival only. A region places work on a lane; it does not sequence the two lanes against each other.

The conservative sequence below uses whole-GM cache maintenance and the soft barrier, so it is safe for multi-cache-line buffers and partial occupancy. sync_ws is an exclusive, zero-initialized 16-element INT32 GM tensor, and participant_count is the total number of participating AIC and AIV cores.

pl.system.cacheinvalid()  # publish all producer cache lines
pl.system.fence()         # wait until they are visible in GM
pl.system.syncall(
    mode=pl.SyncAllMode.SOFT,
    core_type=pl.KernelType.MIX,
    gm_workspace=sync_ws,
    used_cores=participant_count,
)                         # synchronize arrival only
pl.system.cacheinvalid()  # consumer invalidates before reading

For a finer handoff than a whole-barrier rendezvous, pl.system.sync_set / pl.system.sync_wait raise and await a single cross-core event. In a mixed InCore kernel, pin each one to the lane that must run it with core_type=pl.KernelType.AIC or core_type=pl.KernelType.AIV; in an explicitly typed AIC or AIV kernel the lane is already known, so omit the argument.

pl.system.sync_set(0, pipe=pl.PipeType.MTE3, core_type=pl.KernelType.AIV)   # raised on AIV
pl.system.sync_wait(0, pipe=pl.PipeType.MTE2, core_type=pl.KernelType.AIC)    # awaited on AIC

Edge Cases

Fatal pitfall: pl.spmd is an assertion, not a request. You are telling the compiler the blocks are independent. If they are not, the result is a race — not a diagnostic.

Symptom Likely cause Fix
Misplaced tensor op ... should be inside InCore block Operators sit directly in the @pl.jit body Wrap them in with pl.at(level=pl.Level.CORE_GROUP):
with pl.spmd(n): body rejected It neither reads the block index nor dispatches a kernel Read pl.tile.get_block_idx(), or call a kernel
Most pl.write stores vanish, a different set each run Concurrent instances write different elements of one 64-byte cache line — the line, not the element, is what reaches DDR Give each instance whole 64-byte lines, or write from pl.spmd(1); see Memory
optimizations= rejected Built up in a variable — the parser reads the AST Write the list inline at the call site
Printed IR cannot be reparsed A device-size query was bound to a name before use Write the call inline where it is used
vector op '...' sits outside every pl.split_aiv region The function opens a region, so the regions own vector placement Wrap that phase in for _ in pl.split_aiv(2, mode=pl.SplitMode.NONE):
cube op '...' inside a pl.split_aiv region A region body is AIV work Move the pl.matmul out of the region
'x' is produced on the CUBE lane ... reads it on the VECTOR lane inside one An unnamed C->V crossing into a region Read it as pl.aiv_shard(x) at the top of the region
'x' is defined inside a pl.split_aiv region but ... reads it on the CUBE lane outside An unnamed V->C crossing out of a region Gather it inside the region: x = pl.aic_gather(x)
'tile.aiv_shard' operand is in Vec, but it transfers a cube-produced value (or '... operand ... is produced on the VECTOR lane by 'tile.full'') pl.aiv_shard of a value the AIV lane produced (pl.full / pl.load) — it already lives on the vector lane, so there is no crossing to name Drop the pl.aiv_shard: author the value at the per-lane extent inside the region, or lane-localize the load with the region's aiv_id
'pl.aiv_shard' crosses the AIC/AIV boundary under ... but ... earlier in this function crosses it under ... One function's crossings mix mode=NONE with a split mode; they share one cross-core pipe Make every crossing agree on split-vs-no-split, drop the crossing from one region, or split the phases into separate pl.at(level=pl.Level.CORE_GROUP) scopes
The cube reads one lane's value at random A V->C crossing out of a mode=NONE region — both lanes push, one shared slot, no arbitration, not diagnosed Gather only a lane-uniform value; use a data-parallel region if the lanes hold different halves
A peer's signal counter reads twice what it should Both AIV lanes ran the same pld.system.notifynot diagnosed Shard the notify by aiv_id, or guard it with if aiv_id == 0:
A rank reads stale data after its pld.system.wait returns Either the double-notify above, or an incomplete cache-publication/fence/barrier/invalidation sequence between the cube and vector phases Shard the notify; add the full GM handoff sequence

See Also