Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/en/dev/language/00-python_syntax.md
Original file line number Diff line number Diff line change
Expand Up @@ -326,9 +326,9 @@ See [Language Guide](../../user/01-language_guide.md#incore-scopes) for examples

`pl.spmd(N)` dispatches a kernel across `N` blocks. Forms:

- `with pl.spmd(N): kernel(...)` — body must be a single call to a pre-defined InCore kernel.
- `with pl.spmd(N): ...` — body is **either** a single call to a pre-defined InCore kernel (direct dispatch, `SpmdScopeStmt(body=Call)`, no inner InCore wrapper) **or** an inline multi-statement block that is auto-outlined into a synthetic InCore region (like the for-form, minus the auto-bound index). An inline body must read the per-block index via `pl.tile.get_block_idx()` — without it every block runs identical work, so the parser rejects it. Captures no producer TaskId.
- `for i in pl.spmd(N): ...` — loop variable binds the per-block index (`pl.tile.get_block_idx()`); the body is auto-outlined into a synthetic InCore region.
- `with pl.spmd(N, deps=[...]) as tid: ...` — **capture form**: mirrors `with pl.at(...) as tid:`. Captures the dispatch's grid-wide producer `pl.Scalar[pl.TASK_ID]` in `tid` (usable as a `deps=` edge, stored into a `pl.array.create(N, pl.TASK_ID)`, or crossed into `pl.manual_scope`), and accepts an inline multi-statement body like the for-form (read the per-block index via `pl.tile.get_block_idx()`). Lowers to an `ir.Submit` whose trailing tuple element is the grid TaskId; `core_num` / `sync_start` ride on the outlined `Spmd` Function attrs. See [Manual dependency primitives](#manual-dependency-primitives).
- `with pl.spmd(N, deps=[...]) as tid: ...` — **capture form**: mirrors `with pl.at(...) as tid:`. Same body shapes as the plain form above, and additionally captures the dispatch's grid-wide producer `pl.Scalar[pl.TASK_ID]` in `tid` (usable as a `deps=` edge, stored into a `pl.array.create(N, pl.TASK_ID)`, or crossed into `pl.manual_scope`). TaskId capture is orthogonal to the inline body — it is the only thing this form adds over the plain form. Lowers to an `ir.Submit` whose trailing tuple element is the grid TaskId; `core_num` / `sync_start` ride on the outlined `Spmd` Function attrs. See [Manual dependency primitives](#manual-dependency-primitives).
- `out, tid = pl.spmd_submit(kernel, *args, core_num=N)` — **submit form**: dispatches the kernel across `N` blocks *and* captures the dispatch's producer `pl.Scalar[pl.TASK_ID]` (the `pl.submit` sibling for a pre-defined kernel). See [Manual dependency primitives](#manual-dependency-primitives).

All three `pl.spmd(...)` scope forms also accept `allow_early_resolve=True` (a boolean literal; same early-dispatch opt-in as `pl.submit` / `pl.at`). It forces the dispatch to lower to an `ir.Submit` even without `as tid` and lowers to `Arg::set_allow_early_resolve(true)`. Rejected on a `pl.cluster()`-nested `pl.spmd` (such a scope is unwrapped into the Group function and never produces a Submit, so the hint would be lost).
Expand Down
4 changes: 2 additions & 2 deletions docs/zh-cn/dev/language/00-python_syntax.md
Original file line number Diff line number Diff line change
Expand Up @@ -322,9 +322,9 @@ for (x,) in pl.while_(init_values=(x_init,)):

`pl.spmd(N)` 把一个 kernel 派发到 `N` 个 block。形式:

- `with pl.spmd(N): kernel(...)` —— body 必须是对一个已声明 InCore kernel 的单次调用。
- `with pl.spmd(N): ...` —— body **既可以**是对一个已声明 InCore kernel 的单次调用(直接派发,`SpmdScopeStmt(body=Call)`,无内层 InCore 包裹),**也可以**是一段内联多语句块,自动外包成一段隐式 InCore 区域(与 for-form 相同,只是不自动绑定索引)。内联 body 必须通过 `pl.tile.get_block_idx()` 读取每个 block 的索引——否则所有 block 执行完全相同的工作,parser 会拒绝。不捕获 producer TaskId
- `for i in pl.spmd(N): ...` —— 循环变量绑定到每个 block 的索引(`pl.tile.get_block_idx()`);body 自动外包成一段隐式 InCore 区域。
- `with pl.spmd(N, deps=[...]) as tid: ...` —— **捕获形式**:与 `with pl.at(...) as tid:` 对称。`tid` 中捕获该分发的 grid 级 producer `pl.Scalar[pl.TASK_ID]`(可用作 `deps=` 边、存入 `pl.array.create(N, pl.TASK_ID)`、或跨入 `pl.manual_scope`),并像 for-form 一样接受内联多语句 body(在 body 内通过 `pl.tile.get_block_idx()` 读取每个 block 的索引)。lower 成一个 `ir.Submit`,其尾部 tuple 元素即 grid TaskId;`core_num` / `sync_start` 记录在外包出的 `Spmd` Function attrs 上。参见下文“手动依赖原语”小节。
- `with pl.spmd(N, deps=[...]) as tid: ...` —— **捕获形式**:与 `with pl.at(...) as tid:` 对称。body 形态与上面的普通形式相同,并额外在 `tid` 中捕获该分发的 grid 级 producer `pl.Scalar[pl.TASK_ID]`(可用作 `deps=` 边、存入 `pl.array.create(N, pl.TASK_ID)`、或跨入 `pl.manual_scope`)。TaskId 捕获与内联 body 正交——这是该形式相比普通形式唯一多出来的能力。lower 成一个 `ir.Submit`,其尾部 tuple 元素即 grid TaskId;`core_num` / `sync_start` 记录在外包出的 `Spmd` Function attrs 上。参见下文“手动依赖原语”小节。
- `out, tid = pl.spmd_submit(kernel, *args, core_num=N)` —— **submit 形式**:将 kernel 在 `N` 个 block 上分发,同时捕获该分发的 producer `pl.Scalar[pl.TASK_ID]`(针对已声明 kernel 的 `pl.submit` 版本)。参见下文“手动依赖原语”小节。

以上三种形式也都接受 `allow_early_resolve=True`(布尔字面量;与 `pl.submit` / `pl.at` 相同的 early-dispatch 选项)。即使不写 `as tid` 也会强制走 `ir.Submit` 形态,并 lower 为 `Arg::set_allow_early_resolve(true)`。在嵌套于 `pl.cluster()` 内的 `pl.spmd` 上会被拒绝(此类 scope 会被 unwrap 进 Group 函数、永远不会产生 Submit,提示会丢失)。
Expand Down
32 changes: 22 additions & 10 deletions python/pypto/language/dsl_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -711,22 +711,26 @@ def spmd(

Usage forms:

1. ``with pl.spmd(n):`` — body must be a single call to a pre-defined
InCore kernel. Can stand alone (implicit cluster) or nest inside
``pl.cluster()``.
1. ``with pl.spmd(n):`` — body is either a single call to a pre-defined
InCore kernel (direct dispatch), or an inline multi-statement block that
is auto-outlined into a synthetic InCore kernel (like the loop form, minus
the auto-bound index). An inline body must read the per-block index via
``pl.tile.get_block_idx()`` — without it every block runs identical work.
Captures no producer TaskId (use form 3 for that). Can stand alone
(implicit cluster) or nest inside ``pl.cluster()``.

2. ``for i in pl.spmd(n):`` — loop-style. The iteration variable binds
the per-block index (equivalent to ``pl.tile.get_block_idx()``); the
body is auto-outlined into a synthetic InCore function, so inline
tile/tensor ops work without a separate ``@pl.function(type=InCore)``
declaration.

3. ``with pl.spmd(n, deps=[...]) as tid:`` — captures the grid dispatch's
producer ``Scalar[TASK_ID]`` in ``tid`` (mirroring ``with pl.at(...) as
tid:``), usable as a ``deps=`` edge on later tasks, stored into a
``pl.array.create(N, pl.TASK_ID)``, or crossing into ``pl.manual_scope``.
Unlike form 1, this form accepts an inline multi-statement body (like the
loop form); read the per-block index inside via ``pl.tile.get_block_idx()``.
3. ``with pl.spmd(n, deps=[...]) as tid:`` — same body shapes as form 1, and
additionally captures the grid dispatch's producer ``Scalar[TASK_ID]`` in
``tid`` (mirroring ``with pl.at(...) as tid:``), usable as a ``deps=`` edge
on later tasks, stored into a ``pl.array.create(N, pl.TASK_ID)``, or
crossing into ``pl.manual_scope``. TaskId capture is the only thing this
adds over form 1 — it is orthogonal to the inline body.

Optional ``optimizations=[pl.split(mode)]`` applies to the inner InCore scope
(auto-generated for the for-form and the ``as tid`` form, wrapped around the
Expand Down Expand Up @@ -761,10 +765,18 @@ def spmd(
Context manager / loop iterator for the SPMD scope.

Examples:
>>> # Single-kernel context-manager form
>>> # Single-kernel context-manager form (direct dispatch)
>>> with pl.spmd(4):
... out = self.kernel(a, b, out)
>>>
>>> # Inline context-manager form — no TaskId; read the block index yourself
>>> with pl.spmd(4):
... i = pl.tile.get_block_idx()
... offset = i * 128
... tile_a = pl.load(a, [offset, 0], [128, 128])
... tile_b = pl.load(b, [offset, 0], [128, 128])
... out = pl.store(pl.add(tile_a, tile_b), [offset, 0], out)
>>>
>>> # Loop form — body runs per-block with i = tile.get_block_idx()
>>> for i in pl.spmd(4):
... offset = i * 128
Expand Down
Loading
Loading