diff --git a/docs/en/dev/language/00-python_syntax.md b/docs/en/dev/language/00-python_syntax.md index 6f0e6b417..da9c0a1de 100644 --- a/docs/en/dev/language/00-python_syntax.md +++ b/docs/en/dev/language/00-python_syntax.md @@ -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). diff --git a/docs/zh-cn/dev/language/00-python_syntax.md b/docs/zh-cn/dev/language/00-python_syntax.md index a88c61927..543e41a67 100644 --- a/docs/zh-cn/dev/language/00-python_syntax.md +++ b/docs/zh-cn/dev/language/00-python_syntax.md @@ -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,提示会丢失)。 diff --git a/python/pypto/language/dsl_api.py b/python/pypto/language/dsl_api.py index 6420e3284..98c732da8 100644 --- a/python/pypto/language/dsl_api.py +++ b/python/pypto/language/dsl_api.py @@ -711,9 +711,13 @@ 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 @@ -721,12 +725,12 @@ def spmd( 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 @@ -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 diff --git a/python/pypto/language/parser/ast_parser.py b/python/pypto/language/parser/ast_parser.py index 3d6e8243e..e1b1e4c85 100644 --- a/python/pypto/language/parser/ast_parser.py +++ b/python/pypto/language/parser/ast_parser.py @@ -3587,6 +3587,128 @@ def _reject_spmd_early_resolve_in_cluster(self, allow_early_resolve: bool, span: "cluster) to keep the hint.", ) + @staticmethod + def _spmd_body_reads_block_idx(body: "list[ast.stmt]") -> bool: + """True if any statement in an inline SPMD body calls ``get_block_idx()``. + + An inline (auto-outlined) ``pl.spmd`` body distinguishes blocks solely via + the per-block index; without it every block executes identical work — almost + always a bug, and the reason the body is being outlined into a per-block + kernel at all. The single-call direct-dispatch shape is exempt (the callee + reads the index internally), so this is only consulted for inline bodies. + + Matched at the AST layer (no IR ``Op`` exists yet) by the trailing call name, + so every valid spelling of the API counts regardless of receiver: + ``pl.get_block_idx()`` (the top-level alias real models use), the qualified + ``pl.tile.get_block_idx()`` / ``tile.get_block_idx()``, and a bare + ``get_block_idx()`` imported directly. Matching by name only is deliberately + lenient: ``get_block_idx`` is unique to this API (no other DSL object exposes + it), and being lenient here is far safer than rejecting a real body that + distinguishes blocks. ``ast.walk`` recurses the whole body subtree, so a + nested use (inside a ``pl.range`` loop or an expression argument) is found. + """ + for body_stmt in body: + for node in ast.walk(body_stmt): + if isinstance(node, ast.Call): + func = node.func + if (isinstance(func, ast.Attribute) and func.attr == "get_block_idx") or ( + isinstance(func, ast.Name) and func.id == "get_block_idx" + ): + return True + return False + + def _emit_spmd_body( # noqa: PLR0913 — args map 1:1 to the SpmdScopeStmt fields + self, + stmt: ast.With, + span: "ir.Span", + scope_kind: "ir.ScopeKind", + core_num: "ir.Expr", + sync_start: bool, + name_hint: str, + split_mode: "ir.SplitMode | None", + split_slot_num: "int | None", + scope_attrs: "list[tuple[str, Any]]", + ) -> None: + """Emit the ``SpmdScopeStmt`` body shared by the plain and ``as tid`` with-forms. + + The two forms differ only in ``scope_attrs`` (the ``as tid`` form adds + ``task_id_var`` / ``manual_dep_edges``); the body dispatch is identical: + + * single call + no split → ``SpmdScopeStmt(body=Call)`` with no inner InCore + wrapper — the historical direct-dispatch shape (the callee is a pre-defined + kernel that reads the block index internally). This is also the shape + ``OutlineIncoreScopes`` leaves behind once an inline body is outlined, so + the IR round-trips identically across passes. + * inline multi-statement body, or single-call + split → wrap in + ``InCoreScopeStmt(split,
)`` for ``OutlineIncoreScopes`` to outline + into a synthetic per-block kernel, exactly like ``for i in pl.spmd(n):``. + Such an inline body must read the per-block index (see below). + """ + # A single body statement whose value is a Call — Assign/AnnAssign/Expr all + # expose a `.value`, so one membership test covers the three call-carrying + # statement kinds (`x = f()`, `x: T = f()`, and a bare `f()`). + body_stmt = stmt.body[0] if len(stmt.body) == 1 else None + is_single_call = isinstance(body_stmt, (ast.Assign, ast.AnnAssign, ast.Expr)) and isinstance( + body_stmt.value, ast.Call + ) + # An inline (auto-outlined) body must read the per-block index — the + # single-call dispatch is exempt (its callee reads it internally). Unlike + # the for-form, the with-forms do not bind the index for you, so require an + # explicit ``pl.tile.get_block_idx()`` somewhere in the body. + if not is_single_call and not self._spmd_body_reads_block_idx(stmt.body): + raise ParserSyntaxError( + "inline `with pl.spmd(...)` body must read the per-block index via " + "`pl.tile.get_block_idx()`; without it every block runs identical work.", + span=span, + hint="Add `i = pl.tile.get_block_idx()` inside the scope, or use " + "`for i in pl.spmd(n):` to bind the block index automatically.", + ) + if is_single_call and split_mode is None: + # Historical no-InCore-wrapper shape. Any ``scope_attrs`` + # (allow_early_resolve, and for the ``as tid`` form task_id_var / + # manual_dep_edges) ride on the SpmdScopeStmt. + self._parse_scope_body( + stmt, + scope_kind, + span, + name_hint=name_hint, + core_num=core_num, + sync_start=sync_start, + attrs=scope_attrs or None, + ) + return + # split= hint or an inline multi-statement body requires an inner + # InCoreScopeStmt to carry split_ / be outlined. Build the scope directly + # instead of routing through _parse_scope_body, so merge any forward-sticky + # pl.dump_tag tensors onto it here (see _parse_spmd_for_loop for the full + # rationale). + spmd_name_hint, incore_name_hint = _split_spmd_for_loop_name_hints(name_hint) + incore_attrs = self._merge_forward_sticky_dump(None, ir.ScopeKind.InCore) + incore_attrs = self._append_split_slot_num_attr(incore_attrs, split_slot_num) + with self.builder.scope( + scope_kind, + span, + name_hint=spmd_name_hint, + core_num=core_num, + sync_start=sync_start, + attrs=scope_attrs or None, + ): + with self._scope_kind_context(scope_kind): + self.scope_manager.enter_scope("spmd_with") + with self.builder.scope( + ir.ScopeKind.InCore, + span, + split=split_mode, + name_hint=incore_name_hint, + attrs=incore_attrs, + ): + with self._scope_kind_context(ir.ScopeKind.InCore): + self.scope_manager.enter_scope("spmd_with_incore") + self._parse_body_siblings(stmt.body) + self._discard_tail_block_comments(stmt.body, upper_line=stmt.end_lineno) + self.scope_manager.exit_scope(leak_vars=True) + self.scope_manager.exit_scope(leak_vars=True) + def _parse_spmd_scope( self, stmt: ast.With, @@ -3596,19 +3718,26 @@ def _parse_spmd_scope( ) -> None: """Parse ``with pl.spmd(...):`` / ``with pl.spmd(...) as tid:`` into a ScopeStmt(Spmd). - Two forms: + Two forms, differing only in whether the grid dispatch's producer TaskId is + captured — the body shape is identical (see :meth:`_emit_spmd_body`): - * ``with pl.spmd(n): self.kernel(...)`` — wraps a single kernel call - (historical shape; no producer TaskId captured, no ``deps=``). - * ``with pl.spmd(n, deps=[...]) as tid:`` — captures the grid dispatch's - producer ``Scalar[TASK_ID]`` (mirrors ``with pl.at(...) as tid:``) and - accepts an inline multi-statement body that is auto-outlined into an - InCore kernel, exactly like ``for i in pl.spmd(n):``. The per-block - index is read inside the body via ``pl.tile.get_block_idx()``. + * ``with pl.spmd(n):`` — no captured TaskId, no ``deps=``. Accepts either a + single kernel call (historical direct-dispatch shape) or an inline + multi-statement body auto-outlined into an InCore kernel (like + ``for i in pl.spmd(n):``, minus the auto-bound loop var — read the + per-block index inside via ``pl.tile.get_block_idx()``). + * ``with pl.spmd(n, deps=[...]) as tid:`` — same body shapes, and + additionally captures the producer ``Scalar[TASK_ID]`` (mirrors + ``with pl.at(...) as tid:``) so it can feed a ``deps=`` edge. + + TaskId capture and inline bodies are orthogonal: the inline body is outlined + the same way with or without ``as tid``; ``as tid`` only adds the + ``task_id_var`` attr that makes the dispatch lower to an ``ir.Submit``. """ with_hint = ( - "Use 'with pl.spmd(4):' with a single function call inside, or " - "'with pl.spmd(4) as tid:' to capture the dispatch TaskId." + "Use 'with pl.spmd(4):' with a single call or an inline block that reads " + "'pl.tile.get_block_idx()', or 'with pl.spmd(4) as tid:' to also capture " + "the dispatch TaskId." ) # ``deps=`` is accepted ONLY with ``as tid`` — gate it by keyword presence, # not by the resolved list being non-empty. _parse_submit_deps_kwarg @@ -3656,81 +3785,22 @@ def _parse_spmd_scope( self._reject_spmd_early_resolve_in_cluster(allow_early_resolve, span) spmd_attrs: list[tuple[str, Any]] = [("allow_early_resolve", True)] if allow_early_resolve else [] - # No ``as tid``: the historical single-kernel-call with-form. ``deps=`` was - # already rejected above (allow_deps=False), so dep_vars is empty here. - # Validate body is exactly one statement that is a function call. - # The loop form (for i in pl.spmd(n):) and the `as tid` with-form are - # what accept inline multi-statement bodies. - spmd_hint = ( - "The 'with pl.spmd()' form (without 'as tid') wraps a single kernel call. " - "Use 'with pl.spmd(4) as tid:' to write inline tile/tensor ops and capture the " - "dispatch TaskId, or 'for i in pl.spmd(4):' for an inline loop body." - ) - if len(stmt.body) != 1: - raise ParserSyntaxError( - f"pl.spmd() body must contain exactly one statement, got {len(stmt.body)}", - span=self.span_tracker.get_span(stmt), - hint=spmd_hint, - ) - body_stmt = stmt.body[0] - is_call = ( - (isinstance(body_stmt, ast.Assign) and isinstance(body_stmt.value, ast.Call)) - or (isinstance(body_stmt, ast.AnnAssign) and isinstance(body_stmt.value, ast.Call)) - or (isinstance(body_stmt, ast.Expr) and isinstance(body_stmt.value, ast.Call)) + # No ``as tid``: the plain with-form. ``deps=`` was already rejected above + # (allow_deps=False), so dep_vars is empty here. The shared helper keeps the + # historical single-call direct-dispatch shape and outlines an inline + # multi-statement body into a synthetic InCore kernel — identical to the + # ``as tid`` form, minus the captured TaskId. + self._emit_spmd_body( + stmt, + span, + scope_kind, + core_num, + sync_start, + name_hint, + split_mode, + split_slot_num, + spmd_attrs, ) - if not is_call: - raise ParserSyntaxError( - "pl.spmd() body statement must be a function call", - span=self.span_tracker.get_span(stmt), - hint=spmd_hint, - ) - if split_mode is None: - # No optimizations — preserve the historical IR shape: - # SpmdScopeStmt(