feat(passes): materialize distributed tensor comm ctx

[hashiqiqixian]

Summary

• Add MaterializeDistTensorCtx to materialize explicit CommCtxType params and args for DistributedTensorType function params.
• Wire the new pass into the default/debug pipelines after LowerHostTensorCollectives and before the final Simplify.
• Update PTO, orchestration, distributed, and backend codegen to consume explicit CommCtx values instead of implicit trailing ctx synthesis.
• Add collision-safe ctx naming for existing params/locals such as data_ctx.
• Update C++ registration, Python bindings, type stubs, docs, and pass-order tests.

Testing

• git diff --check
• python -m py_compile ...
• keyword scan for old trailing ctx / ctx-name guessing paths

Note: full pytest was not run because pytest is not installed in this environment.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ffcb6a59-6773-4f5f-9045-ff3d942a8983

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Introduces a new IR transformation pass, MaterializeDistTensorCtx, which explicitly threads CommCtxType parameters/arguments through function signatures and call/submit sites for DistributedTensorType parameters. Updates PTO, distributed, and orchestration codegen backends to consume the explicit context, wires the pass into pipelines and Python bindings, and updates documentation and tests.

Changes

MaterializeDistTensorCtx pass and codegen integration

Layer / File(s)	Summary
Pass contract, factory, and bindings include/pypto/ir/transforms/pass_properties.h, include/pypto/ir/transforms/passes.h, python/bindings/modules/passes.cpp, python/pypto/pypto_core/passes.pyi, python/pypto/ir/pass_manager.py, CMakeLists.txt	Adds PassProperties entry, C++/Python factory declarations, Python stub with docstring and __all__ export, and pipeline registration; adds new source file to the build.
Pass implementation src/ir/transforms/materialize_dist_tensor_ctx_pass.cpp	Implements ctx-variable planning per function, signature extension, and Call/Submit rewriting with prefix pld.system.get_comm_ctx assignments and updated ArgDirection.
PTO codegen consumption include/pypto/codegen/pto/pto_codegen.h, src/codegen/pto/pto_codegen.cpp	PTOCodegen now emits CommCtxType params as !pto.ptr<i64> already present in func.func signatures, pairs them with DistributedTensor params, and stops appending trailing ctx pointers.
Distributed codegen consumption include/pypto/codegen/distributed/distributed_codegen.h, src/codegen/distributed/distributed_codegen.cpp	Adds ResolveCommCtxArg helper, scalar-like binding for CommCtxType params, special-case handling of get_comm_ctx assignments, and TaskArgs.add_scalar lowering for CommCtx call args.
Orchestration codegen wrapper plumbing src/codegen/orchestration/orchestration_codegen.cpp	Refactors BuildWrapperReorderedParams to a unified ParamEntry vector, adds get_comm_ctx assignment lowering, updates GetCppType for CommCtx→uint64_t, and updates Spmd/Group call generation.
Backend kernel unpacking python/pypto/backend/pto_backend.py	Unpacks CommCtxType parameters as scalar-like __gm__ int64_t* pointers alongside scalars, and restricts *_conv name bookkeeping to ScalarType.
Tests tests/ut/ir/transforms/test_materialize_dist_tensor_ctx.py, tests/ut/ir/transforms/test_pass_manager.py, tests/ut/codegen/distributed/*, tests/ut/codegen/test_spmd_scope_tid_codegen.py	Adds new pass test module and updates codegen test pipelines/assertions to include and verify materialize_dist_tensor_ctx().
Documentation docs/en/dev/passes/*, docs/zh-cn/dev/passes/*, docs/en/dev/codegen/..., docs/zh-cn/dev/codegen/..., docs/en/dev/language/..., docs/zh-cn/dev/language/..., docs/en/user/..., docs/zh-cn/user/..., .claude/rules/pass-doc-ordering.md	Adds new pass documentation pages and updates pass ordering tables/links across renumbered pass docs.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• hw-native-sys/pypto#1676: Both PRs modify orchestration codegen's Spmd/Group wrapper parameter plumbing.
• hw-native-sys/pypto#1685: Both PRs adjust DistributedTensor↔CommContext plumbing in PTO codegen.
• hw-native-sys/pypto#1917: Both PRs modify orchestration codegen wrapper/call parameter reordering and CommCtx scalar threading with related spmd/Group tests.

Suggested labels: enhancement

Poem

A tensor hops across the divide,
with CommCtx now riding by its side.
No more hidden trailing pointer trick—
the IR spells it out, clean and quick.
Hop, pass, hop—signatures grow,
and every rank knows just where to go! 🐇✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: materializing distributed tensor communication contexts.
Description check	✅ Passed	The description is directly aligned with the changeset and accurately summarizes the new pass, pipeline wiring, and codegen updates.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces the MaterializeDistTensorCtx pass, which materializes an explicit CommCtxType parameter and argument for each DistributedTensorType function parameter, making the communication context flow explicit in the IR. It updates the pipeline, codegen paths, and Python bindings accordingly, and adds comprehensive unit tests. However, a critical issue was identified in src/codegen/orchestration/orchestration_codegen.cpp, where widespread character encoding corruption occurred, introducing a Byte Order Mark (BOM) and corrupting various UTF-8 characters into garbage text.

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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b2b692c2a8

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/ut/codegen/test_spmd_scope_tid_codegen.py (1)

39-55: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Update the stale ext_signal_ctx expectations in the two existing DistributedTensor tests.
_codegen() now materializes dist-tensor ctx for every test in this class, so these cases should emit signal_ctx directly, matching the new test. The assertions at lines 520 and 575 still look like pre-pass output and likely need to be changed to params_t0.add_scalar(signal_ctx);.

🤖 Prompt for 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.

In `@tests/ut/codegen/test_spmd_scope_tid_codegen.py` around lines 39 - 55, The
`DistributedTensor` test expectations are stale because `_codegen()` now runs
`materialize_dist_tensor_ctx()` before codegen, so the generated orchestration
output should use the materialized `signal_ctx` form instead of the older
`ext_signal_ctx`-style output. Update the two assertions in the existing
`DistributedTensor` tests to match the new output from
`codegen.generate_orchestration(...)`, specifically the
`params_t0.add_scalar(signal_ctx);` form, so they align with the behavior
exercised by `_codegen()`.

🤖 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 `@src/codegen/orchestration/orchestration_codegen.cpp`:
- Line 1: Remove the UTF-8 BOM from the start of orchestration_codegen.cpp so
the file begins directly with the header comment instead of the EF BB BF marker.
Keep the rest of the file unchanged; just save the file as plain UTF-8 without
BOM.

In `@src/ir/transforms/materialize_dist_tensor_ctx_pass.cpp`:
- Around line 74-90: BuildFunctionCtxPlan currently seeds used_names only from
func->params_, so MakeUniqueName can still generate a ctx parameter that
collides with a name already used in the function body. Update
BuildFunctionCtxPlan to also collect body-local names from the FunctionPtr
before creating ctx vars, and keep using MakeCtxParamBaseName plus
MakeUniqueName so ExtendFunctionSignature receives non-conflicting ctx parameter
names.

---

Outside diff comments:
In `@tests/ut/codegen/test_spmd_scope_tid_codegen.py`:
- Around line 39-55: The `DistributedTensor` test expectations are stale because
`_codegen()` now runs `materialize_dist_tensor_ctx()` before codegen, so the
generated orchestration output should use the materialized `signal_ctx` form
instead of the older `ext_signal_ctx`-style output. Update the two assertions in
the existing `DistributedTensor` tests to match the new output from
`codegen.generate_orchestration(...)`, specifically the
`params_t0.add_scalar(signal_ctx);` form, so they align with the behavior
exercised by `_codegen()`.

🪄 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: 694d8adf-c528-486b-854e-59c132c4c6b8

📥 Commits

Reviewing files that changed from the base of the PR and between 79e21a1 and b2b692c.

📒 Files selected for processing (35)

• .claude/rules/pass-doc-ordering.md
• CMakeLists.txt
• docs/en/dev/codegen/01-orchestration_codegen.md
• docs/en/dev/language/00-python_syntax.md
• docs/en/dev/passes/00-pass_manager.md
• docs/en/dev/passes/37-materialize_comm_domain_scopes.md
• docs/en/dev/passes/38-lower_host_tensor_collectives.md
• docs/en/dev/passes/39-materialize_dist_tensor_ctx.md
• docs/en/dev/passes/40-materialize_runtime_scopes.md
• docs/en/user/01-language_guide.md
• docs/zh-cn/dev/codegen/01-orchestration_codegen.md
• docs/zh-cn/dev/language/00-python_syntax.md
• docs/zh-cn/dev/passes/00-pass_manager.md
• docs/zh-cn/dev/passes/37-materialize_comm_domain_scopes.md
• docs/zh-cn/dev/passes/38-lower_host_tensor_collectives.md
• docs/zh-cn/dev/passes/39-materialize_dist_tensor_ctx.md
• docs/zh-cn/dev/passes/40-materialize_runtime_scopes.md
• docs/zh-cn/user/01-language_guide.md
• include/pypto/codegen/distributed/distributed_codegen.h
• include/pypto/codegen/pto/pto_codegen.h
• include/pypto/ir/transforms/pass_properties.h
• include/pypto/ir/transforms/passes.h
• python/bindings/modules/passes.cpp
• python/pypto/backend/pto_backend.py
• python/pypto/ir/pass_manager.py
• python/pypto/pypto_core/passes.pyi
• src/codegen/distributed/distributed_codegen.cpp
• src/codegen/orchestration/orchestration_codegen.cpp
• src/codegen/pto/pto_codegen.cpp
• src/ir/transforms/materialize_dist_tensor_ctx_pass.cpp
• tests/ut/codegen/distributed/test_distributed_pto_codegen.py
• tests/ut/codegen/distributed/test_host_orch_distributed.py
• tests/ut/codegen/test_spmd_scope_tid_codegen.py
• tests/ut/ir/transforms/test_materialize_dist_tensor_ctx.py
• tests/ut/ir/transforms/test_pass_manager.py