feat(language): decouple pl.spmd inline body from TaskId capture#1947
Conversation
The plain `with pl.spmd(n):` form previously required a single kernel call. Decouple inline-body support from TaskId capture: the plain form now also accepts an inline multi-statement body, auto-outlined into a synthetic InCore kernel exactly like the for-form and the `as tid` form, without requiring `as tid`. TaskId capture (`as tid`) is now purely additive — it only adds the task_id_var attr that lowers the dispatch to an ir.Submit. Factor the shared body-dispatch tail (single-call direct-dispatch vs InCore-wrap for outlining) out of `_parse_spmd_scope` and `_parse_spmd_scope_with_tid` into a new `_emit_spmd_body` helper; both forms now differ only in scope_attrs. Add a block-index requirement for inline bodies: an auto-outlined body must read the per-block index via `pl.tile.get_block_idx()`, otherwise every block runs identical work. Enforced in the shared helper via `_spmd_body_reads_block_idx`, so both with-forms reject it consistently; the single-call direct-dispatch shape is exempt (its callee reads the index internally). Update the spmd() docstring, the EN/zh-CN python-syntax docs, and add parser + JIT end-to-end tests covering the no-tid inline path and the missing-block-index rejection on both with-forms.
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughRefactors Changespl.spmd inline body support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant User as Developer code
participant Parser as ASTParser
participant Helper as _emit_spmd_body
participant IR as IR builder
User->>Parser: with pl.spmd(N): ... (inline body)
Parser->>Helper: emit body(single-call or inline)
Helper->>Helper: check get_block_idx() usage
alt inline body missing get_block_idx()
Helper-->>Parser: raise ParserSyntaxError
else valid body
Helper->>IR: build InCoreScopeStmt or direct call
IR-->>Parser: SpmdScopeStmt
end
Parser-->>User: parsed program
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request decouples inline-body support from TaskId capture in pl.spmd(N). It allows the plain with pl.spmd(N): form (without as tid) to accept an inline multi-statement body that is auto-outlined into a synthetic InCore kernel, while ensuring that any inline body explicitly reads the per-block index via pl.tile.get_block_idx(). The parser has been refactored to share the body-emitting logic, and corresponding documentation, docstrings, and unit tests have been added or updated. As there are no review comments provided, I have no additional feedback to offer.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
There was a problem hiding this comment.
Pull request overview
This PR updates the PyPTO language parser and documentation so the plain with pl.spmd(n): form can host an inline multi-statement SPMD body without forcing as tid TaskId capture, making TaskId capture purely additive.
Changes:
- Factor shared SPMD with-scope body emission into
_emit_spmd_body, enabling inline bodies for the no-as tidform while preserving the historical single-call direct-dispatch shape. - Enforce a new parser rule for auto-outlined inline bodies: they must read
pl.tile.get_block_idx()(applies to both with-forms), and add corresponding parser + JIT end-to-end tests. - Update
spmd()docstring and EN/zh-CN language syntax docs to reflect the decoupled forms and the block-index requirement.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/ut/language/parser/test_scope_parsing.py | Adds parser-level regression tests for the new no-as tid inline with-form and shared block-index enforcement. |
| tests/ut/jit/test_spmd.py | Adds an end-to-end JIT test ensuring the new inline no-as tid with-form compiles and does not lower to Submit. |
| python/pypto/language/parser/ast_parser.py | Implements shared _emit_spmd_body and block-index detection/enforcement for inline with-form SPMD bodies. |
| python/pypto/language/dsl_api.py | Updates pl.spmd() docstring to document the new inline plain with-form and the block-index requirement. |
| docs/en/dev/language/00-python_syntax.md | Updates the English language syntax docs for the new SPMD forms and limitation. |
| docs/zh-cn/dev/language/00-python_syntax.md | Updates the Chinese language syntax docs in sync with the English changes. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 403da27a12
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/pypto/language/parser/ast_parser.py`:
- Around line 3615-3661: The single-call SPMD fast path in _emit_spmd_body is
too broad because it treats any one-line Assign/AnnAssign/Expr call as kernel
dispatch, letting non-kernel calls bypass the block-index check and InCore
wrapping. Narrow is_single_call so it only matches self.<method>(...)
dispatches, and keep all other single-call bodies subject to the existing
_spmd_body_reads_block_idx guard and outlining behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3b09030a-8b52-408d-81bf-0b2bf3a8700d
📒 Files selected for processing (6)
docs/en/dev/language/00-python_syntax.mddocs/zh-cn/dev/language/00-python_syntax.mdpython/pypto/language/dsl_api.pypython/pypto/language/parser/ast_parser.pytests/ut/jit/test_spmd.pytests/ut/language/parser/test_scope_parsing.py
- Tighten the inline-SPMD block-index guard to match only `<tile>.get_block_idx()` (pl.tile.get_block_idx / tile.get_block_idx), not an unrelated `foo.get_block_idx()` — aligns the check with its own error message and docs (Copilot review). - Add a direct unit test for the AST matcher pinning the accepted spellings and the rejected unrelated-receiver case.
…ive-sys#1947 The previous receiver-restricted match (only `<tile>.get_block_idx()`) regressed a real model (pypto-lib qwen3 14B decode), which reads the per-block index via the top-level `pl.get_block_idx()` alias — a valid API exported at pl top-level, equivalent to `pl.tile.get_block_idx()`. Match `get_block_idx` by trailing call name across every valid spelling regardless of receiver: `pl.get_block_idx()`, `pl.tile.get_block_idx()`, `tile.get_block_idx()`, and a bare imported `get_block_idx()`. Matching by name is deliberately lenient — the name is unique to this API, and leniency is far safer than rejecting a real block-distinguishing body. Update the AST-matcher unit test to cover all spellings (including the pl.get_block_idx() regression case) and add an end-to-end parser test mirroring the qwen3 inline-spmd body.
Summary
Previously, the plain
with pl.spmd(n):form required a single kernel call — the only way to write an inline multi-statement SPMD body in awith-style scope waswith pl.spmd(n) as tid:, which forced you to also capture a producerTaskIdeven when you had no use for it. This coupled two orthogonal concerns.This PR decouples them:
as tidrequired.with pl.spmd(n):now accepts an inline multi-statement body, auto-outlined into a synthetic InCore kernel — exactly likefor i in pl.spmd(n):and theas tidform. The dispatch stays a plainCall(notask_id_var, so it never lowers to anir.Submit).as tidis now purely additive. It shares the same body shapes and only adds thetask_id_varattr that lowers the dispatch to anir.Submit._parse_spmd_scopeand_parse_spmd_scope_with_tidinto a new_emit_spmd_bodyhelper; the two forms now differ only inscope_attrs. Existing IR shapes are byte-identical (scope_attrs or Nonecollapses to the prior values on both paths).New limitation: inline bodies must read the per-block index
An auto-outlined inline body must reference
pl.tile.get_block_idx()— without it every block runs identical work, which is almost always a bug. Enforced in the shared helper via_spmd_body_reads_block_idx, so both with-forms reject it consistently. The single-call direct-dispatch shape is exempt (its callee reads the index internally); theforform is unaffected (it auto-binds the index).Example
Testing
TestSpmdInlineWithForm): no-tid inline path wraps InCore with notask_id_var, notask_invalidplaceholder, split wrapping, print→parse round-trip, and missing-block-index rejection on both the plain andas tidforms.test_jit_spmd_inline_with_form_no_tid): the inline plain form compiles through the full pipeline and the dispatch stays aCall(never aSubmit).pl.cluster()-nested inline forms (the latter works becauseOutlineIncoreScopesat pass 08 outlines the inline body beforeOutlineClusterScopesat pass 09).Docs
spmd()docstring (forms 1 & 3 + inline example) and the EN + zh-CNdocs/.../language/00-python_syntax.mdspmd sections updated.