feat(passes): materialize distributed tensor comm ctx

.claude/rules/pass-doc-ordering.md

@@ -52,7 +52,8 @@ Developers read pass docs sequentially to understand the compilation pipeline. I
 | 37 | `37-synthesize_allreduce_signals.md` | 37th pass (distributed: host allreduce optional signal -> explicit internal signal IR) |
 | 38 | `38-materialize_comm_domain_scopes.md` | 38th pass (distributed: WindowBuffer + CommDomainScopeStmt wrappers in each host_orch body; runs immediately before LowerHostTensorCollectives) |
 | 39 | `39-lower_host_tensor_collectives.md` | 39th pass (host-level tensor collectives -> internal builtin chip dispatches; runs after comm-domain scopes) |
-| 40 | `40-materialize_runtime_scopes.md` | Last pass (after the final Simplify; inserts AUTO RuntimeScopeStmt so orchestration codegen emits PTO2_SCOPE 1:1) |
+| 40 | `40-materialize_dist_tensor_ctx.md` | 40th pass (materializes explicit CommCtx params/args for DistributedTensor params; runs before the final Simplify) |
+| 41 | `41-materialize_runtime_scopes.md` | Last pass (after the final Simplify; inserts AUTO RuntimeScopeStmt so orchestration codegen emits PTO2_SCOPE 1:1) |
 | 91 | `91-utility_passes.md` | Not in Default strategy |
 | 99 | `99-verifier.md` | Infrastructure (not a pipeline pass) |
 

CMakeLists.txt

@@ -161,6 +161,7 @@ set(PYPTO_SOURCES
     src/ir/transforms/synthesize_allreduce_signals_pass.cpp
     src/ir/transforms/materialize_comm_domain_scopes_pass.cpp
     src/ir/transforms/lower_host_tensor_collectives_pass.cpp
+    src/ir/transforms/materialize_dist_tensor_ctx_pass.cpp
     src/ir/transforms/optimize_orch_tensors_pass.cpp
     src/ir/transforms/derive_call_directions_pass.cpp
     src/ir/transforms/auto_derive_task_dependencies_pass.cpp

docs/en/dev/codegen/01-orchestration_codegen.md

@@ -107,7 +107,7 @@ const Tensor& tmp = alloc_0.get_ref(0);
 
 All task submission is wrapped in a top-level `PTO2_SCOPE()`. Codegen no longer
 decides scope placement from the `for` / `if` structure: the
-[MaterializeRuntimeScopes](../passes/40-materialize_runtime_scopes.md) pass
+[MaterializeRuntimeScopes](../passes/41-materialize_runtime_scopes.md) pass
 inserts explicit AUTO `RuntimeScopeStmt` nodes (the function body and each
 `for` / `if` body) into the IR, and codegen emits `PTO2_SCOPE` 1:1 from those
 nodes (manual scopes lower to `PTO2_SCOPE(PTO2ScopeMode::MANUAL)`):

docs/en/dev/language/00-python_syntax.md

@@ -318,7 +318,7 @@ for (x,) in pl.while_(init_values=(x_init,)):
 | `with pl.spmd(N)` / `for i in pl.spmd(N)` | `Spmd` (for-form wraps inner `InCore`) | SPMD multi-block dispatch — see [pl.spmd](#plspmd-multi-block-dispatch) |
 | `pl.spmd(N, optimizations=[pl.split(MODE)])` | `Spmd(InCore(split=MODE))` | Split hint applies to the inner InCore (both forms) |
 | `pl.scope(mode=pl.ScopeMode.MANUAL)` / `pl.manual_scope()` | `Runtime(manual=true)` | Orchestrator MANUAL scope — user manages task ordering. Allowed in either `auto_scope` mode (it is a dependency-semantics choice). See [Manual dependency primitives](#manual-dependency-primitives) |
-| `pl.scope()` | `Runtime(manual=false)` | Orchestrator AUTO scope (`PTO2_SCOPE()`). Hand-placing one requires `@pl.function(auto_scope=False)` (in the default `auto_scope=True` the compiler owns AUTO placement). See [MaterializeRuntimeScopes](../passes/40-materialize_runtime_scopes.md) |
+| `pl.scope()` | `Runtime(manual=false)` | Orchestrator AUTO scope (`PTO2_SCOPE()`). Hand-placing one requires `@pl.function(auto_scope=False)` (in the default `auto_scope=True` the compiler owns AUTO placement). See [MaterializeRuntimeScopes](../passes/41-materialize_runtime_scopes.md) |
 
 See [Language Guide](../../user/01-language_guide.md#incore-scopes) for examples.
 

docs/en/dev/passes/00-pass_manager.md

@@ -89,6 +89,7 @@ struct PassProperties {
 | SynthesizeAllReduceSignals | — | — | — |
 | MaterializeCommDomainScopes | — | CommDomainScopesMaterialized | — |
 | LowerHostTensorCollectives | CommDomainScopesMaterialized | CommDomainScopesMaterialized | — |
+| MaterializeDistTensorCtx | CommDomainScopesMaterialized | CommDomainScopesMaterialized | — |
 | Simplify | — | — | — |
 | MaterializeRuntimeScopes | SplitIncoreOrch, CallDirectionsResolved | RuntimeScopesMaterialized | — |
 
@@ -403,8 +404,9 @@ The PTO-oriented tile stage shared by `Default` and `DebugTileOptimization` is:
 26. [`SynthesizeAllReduceSignals`](37-synthesize_allreduce_signals.md) (distributed: host allreduce optional signal -> explicit internal signal IR)
 27. [`MaterializeCommDomainScopes`](38-materialize_comm_domain_scopes.md) (distributed: WindowBuffer + CommDomainScopeStmt wrappers in each host_orch body; no-op for comm-less programs)
 28. [`LowerHostTensorCollectives`](39-lower_host_tensor_collectives.md) (host-level tensor collectives -> internal builtin chip dispatches)
-29. `Simplify`
-30. [`MaterializeRuntimeScopes`](40-materialize_runtime_scopes.md) (inserts AUTO RuntimeScopeStmt so orchestration codegen emits PTO2_SCOPE 1:1)
+29. [`MaterializeDistTensorCtx`](40-materialize_dist_tensor_ctx.md) (explicit CommCtx params/args for DistributedTensor params)
+30. `Simplify`
+31. [`MaterializeRuntimeScopes`](41-materialize_runtime_scopes.md) (inserts AUTO RuntimeScopeStmt so orchestration codegen emits PTO2_SCOPE 1:1)
 
 `DebugTileOptimization` is a debug-only strategy for inspecting this tile stage
 without the tensor-only prefix passes. Use `Default` for normal compilation and

docs/en/dev/passes/38-materialize_comm_domain_scopes.md

@@ -22,7 +22,7 @@ the IR types so downstream codegen has O(1) access.
 ## Position in the pipeline
 
 ```text
-... -> ExpandManualPhaseFence -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> Simplify (final)
+... -> ExpandManualPhaseFence -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> MaterializeDistTensorCtx -> Simplify (final)
 ```
 
 The pass runs near the end of the default pipeline, immediately before

docs/en/dev/passes/39-lower_host_tensor_collectives.md

@@ -16,7 +16,7 @@ use [`LowerCompositeOps`](12-lower_composite_ops.md).
 ## Position in the pipeline
 
 ```text
-... -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> Simplify (final) -> MaterializeRuntimeScopes
+... -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> MaterializeDistTensorCtx -> Simplify (final) -> MaterializeRuntimeScopes
 ```
 
 The final `Simplify` runs after this pass so any generated loop bounds or

docs/en/dev/passes/40-materialize_dist_tensor_ctx.md

@@ -0,0 +1,80 @@
+# MaterializeDistTensorCtx Pass
+
+Materializes one explicit `CommCtxType` parameter and argument for each
+`DistributedTensorType` function parameter.
+
+## Overview
+
+Distributed tensors need a communication context at every dispatch boundary:
+host orchestration passes a per-rank `device_ctx`, L2 orchestration forwards it
+through task args, and L1 PTO codegen uses it to lower `pld.system.rank`,
+`pld.system.nranks`, `notify`, `wait`, `put`, and remote memory ops.
+
+Older codegen paths synthesized those ctx values independently at several
+sites. This pass makes the ctx flow explicit in IR instead:
+
+1. For every function with `DistributedTensorType` parameters, append matching
+   `CommCtxType` parameters at the tail of the signature, in distributed-tensor
+   parameter order. The appended parameters are `ParamDirection::In`.
+2. For every `Call` / `Submit` to such a function, append matching ctx args.
+   If the distributed tensor arg is a caller parameter, forward the caller's
+   materialized ctx parameter. Otherwise, bind
+   `pld.system.get_comm_ctx(dist)` immediately before the call and pass that
+   result.
+3. If call-site `arg_directions` are already resolved, append matching
+   `ArgDirection::Scalar` entries so downstream codegen can keep treating ctx as
+   ordinary scalar task payload.
+
+The pass runs after `LowerHostTensorCollectives` and before the final
+`Simplify`. At that point host window buffers have already been materialized by
+`MaterializeCommDomainScopes`, host tensor collectives have been lowered, and
+there is still time for the final simplify pass to clean up any forwarding
+aliases.
+
+## Why CommCtx Is Different From Dynamic Dims
+
+Dynamic tensor dimensions can be recovered locally from tensor descriptors at
+the wrapper boundary. A communication context cannot: it is real dataflow across
+host -> orchestration -> task payload -> kernel signature. Keeping it in IR
+prevents codegen sites from drifting out of sync.
+
+## API
+
+| C++ | Python | Level |
+| --- | ------ | ----- |
+| `pass::MaterializeDistTensorCtx()` | `passes.materialize_dist_tensor_ctx()` | Program-level |
+
+```python
+from pypto.pypto_core import passes
+
+program = passes.materialize_dist_tensor_ctx()(program)
+```
+
+## Example
+
+Before:
+
+```python
+def chip_orch(self, data: pld.DistributedTensor[[256], pl.FP32]):
+    return self.kernel(data)
+
+def host_orch(self):
+    data = pld.window(buf, [256], dtype=pl.FP32)
+    self.chip_orch(data, device=r)
+```
+
+After:
+
+```python
+def chip_orch(self, data, data_ctx: pld.CommCtxType):
+    return self.kernel(data, data_ctx)
+
+def host_orch(self):
+    data = pld.window(buf, [256], dtype=pl.FP32)
+    data_ctx = pld.system.get_comm_ctx(data)
+    self.chip_orch(data, data_ctx, device=r)
+```
+
+The kernel body does not need to change. Existing `pld.system.get_comm_ctx(data)`
+uses in the body become pure aliases to the explicit ctx parameter during
+codegen.

docs/en/user/01-language_guide.md

@@ -477,7 +477,7 @@ too; entry `True` + inline `False` is legal and just nests hand scopes
 inside compiler AUTO scopes). `.incore` / `.opaque` reject it — they
 outline into separate kernels. It specializes into
 `@pl.function(..., auto_scope=False)` — see the
-[MaterializeRuntimeScopes pass](../dev/passes/40-materialize_runtime_scopes.md)
+[MaterializeRuntimeScopes pass](../dev/passes/41-materialize_runtime_scopes.md)
 for the resulting scope-placement semantics.
 
 ### `@pl.inline`

docs/zh-cn/dev/codegen/01-orchestration_codegen.md

@@ -106,7 +106,7 @@ const Tensor& tmp = alloc_0.get_ref(0);
 ### 阶段 6–8：任务提交与控制流
 
 所有任务提交包裹在顶层 `PTO2_SCOPE()` 中。codegen 不再依据 `for` / `if` 结构
-决定 scope 位置：[MaterializeRuntimeScopes](../passes/40-materialize_runtime_scopes.md)
+决定 scope 位置：[MaterializeRuntimeScopes](../passes/41-materialize_runtime_scopes.md)
 pass 会向 IR 中插入显式的 AUTO `RuntimeScopeStmt` 节点（函数体以及每个
 `for` / `if` 体），codegen 从这些节点 1:1 地 emit `PTO2_SCOPE`（manual scope
 降级为 `PTO2_SCOPE(PTO2ScopeMode::MANUAL)`）：

docs/zh-cn/dev/language/00-python_syntax.md

@@ -316,7 +316,7 @@ for (x,) in pl.while_(init_values=(x_init,)):
 | `with pl.spmd(N)` / `for i in pl.spmd(N)` | `Spmd`（for-form 内嵌 `InCore`） | SPMD 多 block 派发——见 [pl.spmd](#plspmd-多-block-派发) |
 | `pl.spmd(N, optimizations=[pl.split(MODE)])` | `Spmd(InCore(split=MODE))` | split 提示作用于内层 InCore（两种形式均适用） |
 | `pl.scope(mode=pl.ScopeMode.MANUAL)` / `pl.manual_scope()` | `Runtime(manual=true)` | orchestrator 的 MANUAL scope——由用户管理任务排序。两种 `auto_scope` 模式下都可用（它是依赖语义选择）。见[手工依赖原语](#手工依赖原语) |
-| `pl.scope()` | `Runtime(manual=false)` | orchestrator 的 AUTO scope（`PTO2_SCOPE()`）。手写它需要 `@pl.function(auto_scope=False)`（默认 `auto_scope=True` 下由编译器决定 AUTO 放置）。见 [MaterializeRuntimeScopes](../passes/40-materialize_runtime_scopes.md) |
+| `pl.scope()` | `Runtime(manual=false)` | orchestrator 的 AUTO scope（`PTO2_SCOPE()`）。手写它需要 `@pl.function(auto_scope=False)`（默认 `auto_scope=True` 下由编译器决定 AUTO 放置）。见 [MaterializeRuntimeScopes](../passes/41-materialize_runtime_scopes.md) |
 
 #### `pl.spmd` 多 block 派发
 

docs/zh-cn/dev/passes/00-pass_manager.md

@@ -89,6 +89,7 @@ struct PassProperties {
 | SynthesizeAllReduceSignals | — | — | — |
 | MaterializeCommDomainScopes | — | CommDomainScopesMaterialized | — |
 | LowerHostTensorCollectives | CommDomainScopesMaterialized | CommDomainScopesMaterialized | — |
+| MaterializeDistTensorCtx | CommDomainScopesMaterialized | CommDomainScopesMaterialized | — |
 | Simplify | — | — | — |
 | MaterializeRuntimeScopes | SplitIncoreOrch, CallDirectionsResolved | RuntimeScopesMaterialized | — |
 
@@ -403,8 +404,9 @@ with passes.PassContext([passes.VerificationInstrument(passes.VerificationMode.A
 26. [`SynthesizeAllReduceSignals`](37-synthesize_allreduce_signals.md)（分布式：host allreduce optional signal -> explicit internal signal IR）
 27. [`MaterializeCommDomainScopes`](38-materialize_comm_domain_scopes.md)（分布式：构造 WindowBuffer 并写 CommDomainScopeStmt wrappers in each host_orch body；无通信程序为 no-op）
 28. [`LowerHostTensorCollectives`](39-lower_host_tensor_collectives.md)（host-level tensor collectives -> internal builtin chip dispatches）
-29. `Simplify`
-30. [`MaterializeRuntimeScopes`](40-materialize_runtime_scopes.md)（插入 AUTO RuntimeScopeStmt，使 orchestration codegen 1:1 emit PTO2_SCOPE）
+29. [`MaterializeDistTensorCtx`](40-materialize_dist_tensor_ctx.md)（为 DistributedTensor 参数显式物化 CommCtx 参数/实参）
+30. `Simplify`
+31. [`MaterializeRuntimeScopes`](41-materialize_runtime_scopes.md)（插入 AUTO RuntimeScopeStmt，使 orchestration codegen 1:1 emit PTO2_SCOPE）
 
 `DebugTileOptimization` 只是用于排查 PTO tile 阶段的调试策略，会跳过
 tensor-only 前缀 pass。正常编译和非 strategy 专项测试都应优先使用

docs/zh-cn/dev/passes/38-materialize_comm_domain_scopes.md

@@ -21,7 +21,7 @@
 ## 流水线位置
 
 ```text
-... -> ExpandManualPhaseFence -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> Simplify（最终）
+... -> ExpandManualPhaseFence -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> MaterializeDistTensorCtx -> Simplify（最终）
 ```
 
 本 pass 跑在默认 pipeline 的末尾阶段，位于

docs/zh-cn/dev/passes/39-lower_host_tensor_collectives.md

@@ -15,7 +15,7 @@ builtin chip dispatch。它在 [`MaterializeCommDomainScopes`](38-materialize_co
 ## Pipeline 位置
 
 ```text
-... -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> Simplify（最终） -> MaterializeRuntimeScopes
+... -> SynthesizeAllReduceSignals -> MaterializeCommDomainScopes -> LowerHostTensorCollectives -> MaterializeDistTensorCtx -> Simplify（最终） -> MaterializeRuntimeScopes
 ```
 
 最终的 `Simplify` 位于本 pass 之后，用于继续折叠生成的循环边界或常量表达式，

docs/zh-cn/dev/passes/40-materialize_dist_tensor_ctx.md

@@ -0,0 +1,72 @@
+# MaterializeDistTensorCtx Pass
+
+本 pass 为每个 `DistributedTensorType` 函数参数显式物化一个对应的
+`CommCtxType` 参数和实参。
+
+## 概览
+
+DistributedTensor 的通信上下文需要沿完整调用链传递：host orchestration
+提供每个 rank 的 `device_ctx`，L2 orchestration 通过 task args 转发，L1 PTO
+codegen 再用它降低 `pld.system.rank`、`pld.system.nranks`、`notify`、`wait`、
+`put` 和 remote memory ops。
+
+旧路径在多个 codegen 站点分别合成 ctx，容易漏站点或顺序漂移。本 pass 把这条
+数据流放进 IR：
+
+1. 对每个带 `DistributedTensorType` 参数的函数，按分布式 tensor 参数顺序在签名
+   末尾追加 `CommCtxType` 参数，方向为 `ParamDirection::In`。
+2. 对每个调用点追加对应 ctx 实参。若 distributed tensor 是调用者自己的参数，
+   则转发调用者已有的 ctx 参数；否则在调用前插入
+   `pld.system.get_comm_ctx(dist)` 绑定并传递该结果。
+3. 若调用点已有 `arg_directions`，同步追加 `ArgDirection::Scalar`，让后续
+   codegen 把 ctx 当作普通 scalar task payload 处理。
+
+该 pass 位于 `LowerHostTensorCollectives` 之后、最终 `Simplify` 之前。此时
+host window buffer 已由 `MaterializeCommDomainScopes` 填好，host tensor
+collective 也已降低完成，同时后续仍有一次 simplify 可清理转发别名。
+
+## 与 dynamic dim 的区别
+
+dynamic dim 可以在 wrapper 边界从 tensor descriptor 本地恢复；CommCtx 不行。
+CommCtx 是真实的跨层数据流，必须从 host 到 orchestration、task payload、kernel
+signature 一路传递。把它放进 IR 可以避免多个 codegen 站点各自维护隐式规则。
+
+## API
+
+| C++ | Python | 级别 |
+| --- | ------ | ---- |
+| `pass::MaterializeDistTensorCtx()` | `passes.materialize_dist_tensor_ctx()` | Program-level |
+
+```python
+from pypto.pypto_core import passes
+
+program = passes.materialize_dist_tensor_ctx()(program)
+```
+
+## 示例
+
+Before:
+
+```python
+def chip_orch(self, data: pld.DistributedTensor[[256], pl.FP32]):
+    return self.kernel(data)
+
+def host_orch(self):
+    data = pld.window(buf, [256], dtype=pl.FP32)
+    self.chip_orch(data, device=r)
+```
+
+After:
+
+```python
+def chip_orch(self, data, data_ctx: pld.CommCtxType):
+    return self.kernel(data, data_ctx)
+
+def host_orch(self):
+    data = pld.window(buf, [256], dtype=pl.FP32)
+    data_ctx = pld.system.get_comm_ctx(data)
+    self.chip_orch(data, data_ctx, device=r)
+```
+
+kernel body 不需要修改。body 里已有的 `pld.system.get_comm_ctx(data)` 会在
+codegen 阶段作为指向显式 ctx 参数的纯 alias 处理。

docs/zh-cn/user/01-language_guide.md

@@ -459,7 +459,7 @@ orchestrator（`@pl.jit.host`）和 inline 子函数（`@pl.jit.inline`）上
 只是手放 scope 会嵌套在编译器 AUTO scope 之内）。`.incore` / `.opaque`
 仍会拒绝它——它们外提为独立 kernel。它会 specialize 成
 `@pl.function(..., auto_scope=False)`——具体的 scope 放置语义见
-[MaterializeRuntimeScopes pass](../dev/passes/40-materialize_runtime_scopes.md)。
+[MaterializeRuntimeScopes pass](../dev/passes/41-materialize_runtime_scopes.md)。
 
 ### `@pl.inline`
 

include/pypto/codegen/distributed/distributed_codegen.h

@@ -104,6 +104,12 @@ class DistributedCodegen : public CodegenBase {
   /// Resolve a dispatch call's ``device=`` attr to a Python rank expression.
   [[nodiscard]] std::string ResolveRankExpr(const ir::CallPtr& call) const;
 
+  /// Lower a ``CommCtxType`` call argument. get_comm_ctx-derived locals map to
+  /// the runtime device_ctx expression for the current dispatch rank; explicit
+  /// CommCtx params are forwarded as scalar names.
+  [[nodiscard]] std::string ResolveCommCtxArg(const ir::ExprPtr& arg, const std::string& rank_expr,
+                                              const ir::Span& span) const;
+
   /// Return the emitted Python handle variable for the comm domain that owns
   /// ``wb``.
   [[nodiscard]] std::string GetCommDomainHandleVar(const ir::WindowBufferPtr& wb) const;

include/pypto/codegen/pto/pto_codegen.h

@@ -429,16 +429,15 @@ class PTOCodegen : public CodegenBase {
   [[nodiscard]] std::string GetSpmdSubblockIdxArgSSA() const { return fs_.spmd_subblock_idx_arg; }
 
   /**
-   * @brief SSA name of the CommContext pointer arg appended for a
-   * DistributedTensor param.
-   *
-   * N6 distributed codegen appends one ``!pto.ptr<i64>`` arg per
-   * DistributedTensor parameter at the end of the func.func signature
-   * (after explicit tensor/scalar params, before dynamic-shape ``index``
-   * params). The mapping ``dist_tensor_var → ctx_ssa`` lets the
-   * pld.system.get_comm_ctx / pld.tile.remote_load / pld.tensor.put /
-   * pld.system.notify / pld.system.wait codegen
-   * recover the matching context pointer.
+   * @brief SSA name of the materialized CommContext pointer arg for a
+   * DistributedTensor parameter.
+   *
+   * MaterializeDistTensorCtx adds one explicit ``CommCtxType``
+   * parameter per DistributedTensor parameter. PTOCodegen lowers those
+   * params as ``!pto.ptr<i64>`` scalar arguments and records the
+   * ``dist_tensor_var -> ctx_ssa`` mapping so pld.system.get_comm_ctx /
+   * pld.tile.remote_load / pld.tensor.put / pld.system.notify /
+   * pld.system.wait codegen can recover the matching context pointer.
    *
    * @param dist_var DistributedTensor parameter variable.
    * @return SSA name (e.g. ``%arg7``), or empty string if @p dist_var is

include/pypto/ir/transforms/pass_properties.h

@@ -49,6 +49,10 @@ inline const PassProperties kLowerHostTensorCollectivesProperties{
     .required = {IRProperty::CommDomainScopesMaterialized},
     .produced = {IRProperty::CommDomainScopesMaterialized}};
 
+inline const PassProperties kMaterializeDistTensorCtxProperties{
+    .required = {IRProperty::CommDomainScopesMaterialized},
+    .produced = {IRProperty::CommDomainScopesMaterialized}};
+
 // -- MaterializeRuntimeScopes pass (runs last, after the final Simplify) ------
 //    Inserts explicit AUTO RuntimeScopeStmt nodes for the orchestration function
 //    body and for/if bodies so codegen emits PTO2_SCOPE 1:1 from the IR.