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() with no inner InCore wrapper. The optional - # ``allow_early_resolve`` attr rides on the SpmdScopeStmt; the Spmd - # outliner reads it and threads it onto the synthesised Submit. - self._parse_scope_body( - stmt, - scope_kind, - span, - name_hint=name_hint, - core_num=core_num, - sync_start=sync_start, - attrs=spmd_attrs or None, - ) - else: - # split= hint requires an inner InCoreScopeStmt to carry the - # split_ field. Build SpmdScopeStmt(InCoreScopeStmt(split_=mode, )). - spmd_name_hint, incore_name_hint = _split_spmd_for_loop_name_hints(name_hint) - # Like the for-form, this path builds the InCore 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). - 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=spmd_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_with_tid( # noqa: PLR0913 — args map 1:1 to the SpmdScopeStmt + capture self, @@ -3807,55 +3877,22 @@ def _parse_spmd_scope_with_tid( # noqa: PLR0913 — args map 1:1 to the SpmdSco self.builder.assign(tid_var, placeholder_rhs, span=span) self.builder.push_pending_leading_comments(leading) - # A lone kernel call (no split) keeps the no-InCore-wrapper shape — the - # same SpmdScopeStmt(body=Call) that OutlineIncoreScopes leaves behind once - # an inline body is outlined — so the IR round-trips identically across - # passes. Anything else (inline multi-statement body, or single-call with a - # split hint) wraps in an InCoreScopeStmt, exactly like the for-form. - body_stmt = stmt.body[0] if len(stmt.body) == 1 else None - is_single_call = body_stmt is not None and ( - (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)) - ) - if is_single_call and split_mode is None: - self._parse_scope_body( - stmt, - scope_kind, - span, - name_hint=name_hint, - core_num=core_num, - sync_start=sync_start, - attrs=scope_attrs, - ) - return - - 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, + # Emit the body via the shared helper — identical shape to the plain + # with-form, plus the task_id_var / manual_dep_edges built into scope_attrs + # above: a lone call (no split) keeps the no-InCore-wrapper direct-dispatch + # shape; an inline multi-statement body (or single call + split) wraps in an + # InCoreScopeStmt for outlining and must read the per-block index. + self._emit_spmd_body( + stmt, span, - name_hint=spmd_name_hint, - core_num=core_num, - sync_start=sync_start, - attrs=scope_attrs, - ): - 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) + scope_kind, + core_num, + sync_start, + name_hint, + split_mode, + split_slot_num, + scope_attrs, + ) def _parse_spmd_for_loop(self, stmt: ast.For, iter_call: ast.Call) -> None: """Parse ``for i in pl.spmd(N, ...): body`` into diff --git a/tests/ut/jit/test_spmd.py b/tests/ut/jit/test_spmd.py index e2a0087f6..250d2852e 100644 --- a/tests/ut/jit/test_spmd.py +++ b/tests/ut/jit/test_spmd.py @@ -83,6 +83,50 @@ def entry(a: pl.Tensor, b: pl.Tensor, c: pl.Out[pl.Tensor]): ) +def test_jit_spmd_inline_with_form_no_tid(): + """``with pl.spmd(N):`` with an INLINE body (no ``as tid``) compiles end-to-end. + + Exercises the decoupling of inline-body support from TaskId capture: the plain + with-form auto-outlines the inline body into a synthetic InCore kernel — like the + for-form / as-tid form — without capturing a producer TaskId. The dispatch stays + a plain Call (no task_id_var → not lowered to a Submit), unlike the as-tid form. + """ + torch = pytest.importorskip("torch") + + @jit + def entry(a: pl.Tensor, b: pl.Tensor, c: pl.Out[pl.Tensor]): + with pl.spmd(2): + i = pl.tile.get_block_idx() + offset = i * 64 + tile_a = pl.load(a, [offset, 0], [64, 128]) + tile_b = pl.load(b, [offset, 0], [64, 128]) + c = pl.store(pl.add(tile_a, tile_b), [offset, 0], c) + return c + + post = entry.compile_for_test(torch.randn(128, 128), torch.randn(128, 128), torch.empty(128, 128)) + func_types = {f.func_type for f in post.functions.values()} + assert ir.FunctionType.Spmd in func_types, ( + f"expected an Spmd function from the inline `with pl.spmd()` body, got {func_types}" + ) + # No `as tid`, so the dispatch keeps the plain-Call shape (never a Submit). + orch = next(f for f in post.functions.values() if f.func_type == ir.FunctionType.Orchestration) + spmd_names = {f.name for f in post.functions.values() if f.func_type == ir.FunctionType.Spmd} + submits = [] + + def _walk(node): + if isinstance(node, ir.SeqStmts): + for s in node.stmts: + _walk(s) + elif isinstance(node, ir.AssignStmt): + if isinstance(node.value, ir.Submit) and node.value.op.name in spmd_names: + submits.append(node.value) + elif hasattr(node, "body") and node.body is not None: + _walk(node.body) + + _walk(orch.body) + assert not submits, "plain inline with-form (no as tid) must not lower to a Submit" + + def test_jit_inline_helper_spmd_for_loop(): """A @pl.jit.inline helper using ``for i in pl.spmd(N)`` is spliced + dispatched.""" torch = pytest.importorskip("torch") diff --git a/tests/ut/language/parser/test_scope_parsing.py b/tests/ut/language/parser/test_scope_parsing.py index 74a9f2233..b2668a5e9 100644 --- a/tests/ut/language/parser/test_scope_parsing.py +++ b/tests/ut/language/parser/test_scope_parsing.py @@ -9,9 +9,12 @@ """Unit tests for parsing ScopeStmt with pl.at(level=pl.Level.CORE_GROUP): syntax.""" +import ast + import pypto.language as pl import pytest from pypto import ir +from pypto.language.parser.ast_parser import ASTParser from pypto.language.parser.diagnostics.exceptions import ParserSyntaxError from pypto.language.parser.text_parser import parse_program @@ -1397,6 +1400,227 @@ def bad(x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]: return y +class TestSpmdInlineWithForm: + """``with pl.spmd(n):`` (no ``as tid``) with an inline multi-statement body. + + Decouples inline-body support from TaskId capture: the plain with-form now + auto-outlines an inline body into a synthetic InCore kernel — exactly like the + ``as tid`` form and the for-form — WITHOUT capturing a producer TaskId. The two + concerns are orthogonal (TaskId capture is opt-in via ``as tid``), but an inline + body must still read the per-block index via ``pl.tile.get_block_idx()``. + """ + + @staticmethod + def _descendants(node, cls): + found = [] + + def walk(n): + if isinstance(n, cls): + found.append(n) + if isinstance(n, ir.SeqStmts): + for s in n.stmts: + walk(s) + elif hasattr(n, "body") and n.body is not None: + walk(n.body) + + walk(node) + return found + + @classmethod + def _unique(cls, node, klass): + found = cls._descendants(node, klass) + assert len(found) == 1, f"expected exactly one {klass.__name__}, got {len(found)}" + return found[0] + + def test_inline_with_spmd_no_tid_wraps_incore(self): + """An inline body (no ``as tid``) is auto-outlined into an InCore wrapper and + carries NO task_id_var / manual_dep_edges — the TaskId is not captured.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4, name_hint="stage1"): + i = pl.tile.get_block_idx() + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [i * 128, 0], [128, 128]) + out = pl.store(pl.add(t, t), [i * 128, 0], out) + return out + + main_func = list(Prog.functions.values())[0] + spmd = self._unique(main_func.body, ir.SpmdScopeStmt) + # No TaskId captured — the decoupled feature under test. + assert "task_id_var" not in spmd.attrs + assert "manual_dep_edges" not in spmd.attrs + # The inline body is wrapped in an InCoreScopeStmt for outlining (like the + # for-form / as-tid form), not left as a bare Call. + incore = self._unique(spmd.body, ir.InCoreScopeStmt) + body = incore.body + stmts = list(body.stmts) if isinstance(body, ir.SeqStmts) else [body] + # The user-written get_block_idx is the first body stmt (NOT synthesized). + first = stmts[0] + assert isinstance(first, ir.AssignStmt) + assert isinstance(first.value, ir.Call) + assert first.value.op.name == "tile.get_block_idx" + assert len(stmts) > 1, "inline body should carry multiple statements" + + def test_inline_with_spmd_no_placeholder_before_scope(self): + """Unlike the ``as tid`` form, the plain inline form emits NO + ``AssignStmt(tid, system.task_invalid())`` placeholder before the scope.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4): + i = pl.tile.get_block_idx() + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [i * 128, 0], [128, 128]) + out = pl.store(pl.add(t, t), [i * 128, 0], out) + return out + + main_func = list(Prog.functions.values())[0] + placeholders = [ + s + for s in self._descendants(main_func.body, ir.AssignStmt) + if isinstance(s.value, ir.Call) and s.value.op.name == "system.task_invalid" + ] + assert not placeholders, "plain inline form must not emit a task_invalid placeholder" + + def test_inline_with_spmd_split_wraps_incore_with_split(self): + """``optimizations=[pl.split(...)]`` on the inline plain form sets split_ on the + inner InCore wrapper (same as the for-form / as-tid form).""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4, optimizations=[pl.split(pl.SplitMode.UP_DOWN)]): + i = pl.tile.get_block_idx() + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [i * 128, 0], [128, 128]) + out = pl.store(t, [i * 128, 0], out) + return out + + main_func = list(Prog.functions.values())[0] + spmd = self._unique(main_func.body, ir.SpmdScopeStmt) + assert "task_id_var" not in spmd.attrs + incore = self._unique(spmd.body, ir.InCoreScopeStmt) + assert incore.split == ir.SplitMode.UP_DOWN + + def test_inline_with_spmd_round_trip(self): + """The inline plain form survives print -> parse round-trip (no ``as tid``).""" + + @pl.program + class Original: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4, name_hint="stage1"): + i = pl.tile.get_block_idx() + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [i * 128, 0], [128, 128]) + out = pl.store(pl.add(t, t), [i * 128, 0], out) + return out + + printed = Original.as_python() + assert ".spmd(" in printed and " as tid:" not in printed + Reparsed = pl.parse_program(printed) + ir.assert_structural_equal(Original, Reparsed) + + def test_inline_with_spmd_missing_block_idx_rejected(self): + """An inline body that never reads the per-block index is rejected — without + ``get_block_idx()`` every block runs identical work, so it is almost always a + bug. The single-call direct-dispatch form is exempt (see the regression test + ``test_with_spmd_single_call_still_supported``).""" + with pytest.raises(ParserSyntaxError, match="must read the per-block index"): + + @pl.program + class Bad: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4): + # No pl.tile.get_block_idx() anywhere — every block would run + # identical work writing the same output region. + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [0, 0], [128, 128]) + out = pl.store(pl.add(t, t), [0, 0], out) + return out + + def test_inline_as_tid_missing_block_idx_rejected(self): + """The same block-index requirement applies to the ``as tid`` inline form — + the check lives in the shared body-emit path, so both with-forms enforce it.""" + with pytest.raises(ParserSyntaxError, match="must read the per-block index"): + + @pl.program + class Bad: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4) as tid: # noqa: F841 + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [0, 0], [128, 128]) + out = pl.store(pl.add(t, t), [0, 0], out) + return out + + def test_inline_with_spmd_accepts_top_level_get_block_idx(self): + """Regression (qwen3 decode / pypto-lib-model CI): an inline body that reads + the block index via the top-level ``pl.get_block_idx()`` alias (not the + qualified ``pl.tile.get_block_idx()``) is accepted, not rejected by the guard.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def main( + self, + a: pl.Tensor[[512, 128], pl.FP32], + out: pl.Out[pl.Tensor[[512, 128], pl.FP32]], + ) -> pl.Tensor[[512, 128], pl.FP32]: + with pl.spmd(4, name_hint="fa_fused"): + i = pl.get_block_idx() # top-level alias, as real models use + t: pl.Tile[[128, 128], pl.FP32] = pl.load(a, [i * 128, 0], [128, 128]) + out = pl.store(pl.add(t, t), [i * 128, 0], out) + return out + + main_func = list(Prog.functions.values())[0] + spmd = self._unique(main_func.body, ir.SpmdScopeStmt) + # Outlined (InCore wrapper present), not rejected by the block-index guard. + self._unique(spmd.body, ir.InCoreScopeStmt) + + def test_block_idx_guard_matches_every_get_block_idx_spelling(self): + """The block-index guard matches ``get_block_idx()`` by name, across every + valid spelling regardless of receiver. Regression: the top-level + ``pl.get_block_idx()`` alias (used by real models, e.g. qwen3 decode) must be + accepted — a receiver-restricted match wrongly rejected it.""" + reads = ASTParser._spmd_body_reads_block_idx + # Top-level alias (the regression case), qualified forms, and bare import. + assert reads(ast.parse("x = pl.get_block_idx()").body) + assert reads(ast.parse("x = pl.tile.get_block_idx()").body) + assert reads(ast.parse("x = tile.get_block_idx()").body) + assert reads(ast.parse("x = get_block_idx()").body) + # A nested use (inside an expression argument) still counts. + assert reads(ast.parse("t = pl.load(a, [pl.get_block_idx() * 8, 0], [8, 8])").body) + # A body with no block-index read at all is rejected. + assert not reads(ast.parse("x = pl.load(a, [0, 0], [8, 8])").body) + assert not reads(ast.parse("x = foo.get_subblock_idx()").body) + + class TestSpmdAllowEarlyResolve: """``pl.spmd(..., allow_early_resolve=True)`` — speculative early-dispatch hint.