fix(runtime): group full-rank chip dispatches

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

@@ -73,7 +73,11 @@ This allows extensible operation codegen without modifying the core visitor.
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 ```
 
 ### Phase 2–3: Entry Points
@@ -340,7 +344,11 @@ def orch_basic(
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 extern "C" {
 

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

@@ -95,6 +95,22 @@ For every host-orchestration function (`Function::level_ == Level::HOST` and
    set to `"comm_d<n>"` so codegen emits the matching `__comm_d<n>` handle
    variable verbatim.
 
+9. **Mark full-world dispatch loops.** Stamp
+   `attrs["group_next_level_dispatch"] = true` on a loop only when all of the
+   following are proven: its range is exactly
+   `[0, pld.system.world_size())` with unit step, it has no loop-carried state,
+   its iteration contains exactly one unconditional chip-orchestration call
+   (plus any number of pure `tensor.slice` or `pld.tensor.window` views, but no
+   other calls or `Submit` operations), and that
+   dispatch's `device=` expression is the loop induction variable.
+   Distributed codegen uses this explicit contract to build all per-rank
+   `TaskArgs` first and then call `submit_next_level_group` once. This prevents
+   argument-construction time from becoming rank-start skew. Loops containing
+   nested control flow, any `Submit`, any call other than the two permitted
+   pure view forms, multiple dispatches, a partial/static device range, or a
+   different `device=` expression retain ordinary
+   per-dispatch lowering.
+
 ## Sanity checks
 
 The pass raises `pypto::ValueError` (carrying the alloc's span) if:
@@ -130,6 +146,9 @@ After the pass:
 - Chip-orchestration and InCore parameter types remain `nullopt` on
   `window_buffer_`. N7 codegen reads the back-reference at the *host_orch*
   dispatch site and threads the matching `CommContext` pointer explicitly.
+- Every loop carrying `group_next_level_dispatch = true` is a compiler-proven
+  full-world, one-dispatch-per-rank loop. Codegen consumes the attr but never
+  guesses group eligibility from loop syntax.
 
 ## Pass properties
 

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

@@ -73,7 +73,11 @@ REGISTER_ORCHESTRATION_OP("tensor.slice", TensorSliceHandler);
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 ```
 
 ### 阶段 2–3：入口点
@@ -328,7 +332,11 @@ def orch_basic(
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 extern "C" {
 

docs/zh/dev/passes/41-materialize_comm_domain_scopes.md

@@ -79,6 +79,21 @@ alloc / view / dispatch 点在此时仍然可见。放在较晚阶段还能让
    comm-domain scope 则追加 slot，否则新开一个。`CommDomainScopeStmt wrappers in each host_orch body` 最终
    填充该列表。
 
+8. **包裹 scope**。为每个 comm domain 构造一个嵌套的
+   `CommDomainScopeStmt`；先声明的 domain 在外层。`name_hint_` 使用
+   `"comm_d<n>"`，使 codegen 能直接生成对应的 `__comm_d<n>` handle。
+
+9. **标记全卡 dispatch 循环**。只有在编译器能够证明下列条件全部成立时，
+   才给循环写入 `attrs["group_next_level_dispatch"] = true`：范围严格为
+   `[0, pld.system.world_size())`、步长为 1、没有循环携带状态、每次迭代
+   恰好包含一次无条件 chip-orchestration 调用（允许任意数量的纯 `tensor.slice`
+   或 `pld.tensor.window` view，但不允许其它调用或 `Submit` 操作），且该 dispatch 的
+   `device=` 就是循环归纳变量。分布式
+   codegen 依据这个显式契约，先构造完所有 rank 的
+   `TaskArgs`，再调用一次 `submit_next_level_group`，避免参数构造耗时变成
+   rank 启动偏斜。包含嵌套控制流、`Submit`、上述两种纯 view 以外的调用、多个 dispatch、
+   部分/静态设备范围或不同 `device=` 表达式的循环仍保持逐个 dispatch 的原有 lowering。
+
 ## Sanity 校验
 
 下列情况抛 `pypto::ValueError`（携带 alloc 的 span）：
@@ -110,6 +125,9 @@ pass 运行之后：
 - chip-orchestration 与 InCore 的形参类型 `window_buffer_` 仍是 `nullopt`。
   N7 codegen 在 *host_orch* 的 dispatch 处读取反向引用、再为 chip-orch 显式
   下发对应的 `CommContext` 指针。
+- 每个带 `group_next_level_dispatch = true` 的循环，都已经由编译 pass 证明是
+  “全卡、每 rank 一次 dispatch”的安全形态；codegen 只消费该 attr，不会从
+  循环语法自行猜测是否可以 group。
 
 ## Pass 属性
 

include/pypto/codegen/distributed/distributed_codegen.h

@@ -299,6 +299,11 @@ class DistributedCodegen : public CodegenBase {
   std::set<std::string> declared_vars_;
   bool is_worker_context_{false};
   int task_args_counter_{0};  // Counter for generating unique TaskArgs variable names
+  int group_dispatch_counter_{0};
+  bool group_dispatch_active_{false};
+  std::string group_dispatch_args_var_;
+  std::string group_dispatch_workers_var_;
+  std::string group_dispatch_callee_;
 
   // HOST orchestrator alloc-hoisting state. Populated by
   // CollectHostOrchHoistableAllocs() before EmitFunction() runs on the HOST

include/pypto/ir/transforms/passes.h

@@ -232,6 +232,12 @@ Pass SynthesizeAllReduceSignals();
  *     one comm domain, slots in alloc-source order) and wrap the host_orch
  *     body in nested ``CommDomainScopeStmt`` nodes (outer = first declared
  *     domain, inner = last).
+ *  6. Mark a ``for rank in range(world_size)`` loop for grouped next-level
+ *     publication when its iteration contains exactly one unconditional CHIP
+ *     dispatch pinned to ``rank``, optional pure ``tensor.slice`` or
+ *     ``pld.tensor.window`` views, no other calls or ``Submit`` operations,
+ *     and carries no loop state. Distributed codegen consumes this explicit
+ *     attr; it does not infer the pattern.
  *
  * Sanity-checks (``pypto::ValueError`` on failure):
  *  - Every alloc must have at least one ``pld.tensor.window`` materialisation and

include/pypto/ir/transforms/utils/attrs.h

@@ -201,6 +201,13 @@ inline std::vector<std::pair<std::string, std::any>> StripAttr(
 /// peek through such a scope as if it were AUTO (see ``transform_utils::UnwrapAutoScope``).
 inline constexpr const char* kAttrCompilerAutoManualScopeCandidate = "__compiler_auto_manual_scope_candidate";
 
+/// ``bool`` attr on a HOST-orchestrator ``ForStmt`` whose body is one
+/// rank-pinned CHIP-orchestrator dispatch and whose range is exactly
+/// ``[0, pld.system.world_size())``. ``MaterializeCommDomainScopes`` proves and
+/// stamps this fact; distributed codegen consumes it to build every member's
+/// ``TaskArgs`` before publishing the dispatches as one runtime group.
+inline constexpr const char* kGroupNextLevelDispatchAttr = "group_next_level_dispatch";
+
 // ---------------------------------------------------------------------------
 // ForStmt iter_arg carry classification (produced by ``ClassifyIterArgCarry``)
 // ---------------------------------------------------------------------------

python/pypto/runtime/builtins/collectives/all_to_all/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/allgather/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/allreduce_ring/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/barrier/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/broadcast/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/builtins/collectives/reduce_scatter/templates/entry.cpp.in

@@ -13,7 +13,11 @@
 
 #include <cstdint>
 
+#if __has_include("orchestration_api.h")
+#include "orchestration_api.h"
+#else
 #include "pto_orchestration_api.h"
+#endif
 
 namespace {
 

python/pypto/runtime/distributed_runner.py

@@ -946,6 +946,41 @@ def _submit_chip(orch: Any, callable_id: Any, task_args: Any, config: Any, worke
         config.output_prefix = base
 
 
+def _submit_chip_group(
+    orch: Any,
+    callable_id: Any,
+    task_args_list: list[Any],
+    config: Any,
+    workers: list[int | None],
+) -> Any:
+    """Publish one full-rank CHIP dispatch after every member is prepared.
+
+    With DFX disabled, ``submit_next_level_group`` makes the members one DAG
+    node and activates their target workers together. With DFX enabled, retain
+    :func:`_submit_chip`'s per-rank/per-dispatch output directories; one shared
+    ``CallConfig`` cannot represent a distinct prefix for every group member,
+    and profiling already intentionally perturbs dispatch timing.
+    """
+    if len(task_args_list) != len(workers):
+        raise ValueError("workers length must match task_args_list length")
+    if not task_args_list:
+        raise ValueError("grouped CHIP dispatch requires at least one member")
+    resolved_workers = [_resolve_chip_worker(orch, worker) for worker in workers]
+    if len(set(resolved_workers)) != len(resolved_workers):
+        raise ValueError("workers must not contain duplicate CHIP worker ids")
+    if not config.output_prefix:
+        return orch.submit_next_level_group(
+            callable_id,
+            task_args_list,
+            config,
+            workers=resolved_workers,
+        )
+    return [
+        _submit_chip(orch, callable_id, task_args, config, worker)
+        for task_args, worker in zip(task_args_list, resolved_workers)
+    ]
+
+
 def _clear_dfx_dispatch_dirs(dfx_base: Path) -> None:
     """Remove stale ``rank*/d{k}`` dispatch dirs before a fresh DFX run.
 

src/codegen/distributed/distributed_codegen.cpp

@@ -39,6 +39,7 @@
 #include "pypto/ir/program.h"
 #include "pypto/ir/scalar_expr.h"
 #include "pypto/ir/stmt.h"
+#include "pypto/ir/transforms/utils/attrs.h"
 #include "pypto/ir/transforms/utils/transform_utils.h"
 #include "pypto/ir/type.h"
 
@@ -262,12 +263,17 @@ void DistributedCodegen::EmitImports() {
   // ``_submit_chip`` resolves a comm-less dispatch's chip and namespaces the
   // per-dispatch DFX ``output_prefix`` (``<base>/rank{worker}/d{k}``); the
   // namespacing half is a no-op when DFX is off.
-  emitter_.EmitLine("from pypto.runtime.distributed_runner import _submit_chip");
+  emitter_.EmitLine("from pypto.runtime.distributed_runner import _submit_chip, _submit_chip_group");
 }
 
 void DistributedCodegen::EmitFunction(const ir::FunctionPtr& func) {
   declared_vars_.clear();
   task_args_counter_ = 0;
+  group_dispatch_counter_ = 0;
+  group_dispatch_active_ = false;
+  group_dispatch_args_var_.clear();
+  group_dispatch_workers_var_.clear();
+  group_dispatch_callee_.clear();
   current_func_ = func;
 
   bool is_sub_worker = func->role_.has_value() && *func->role_ == ir::Role::SubWorker;
@@ -837,6 +843,19 @@ void DistributedCodegen::VisitStmt_(const ir::ForStmtPtr& op) {
   std::string step = current_expr_value_;
   current_expr_value_ = "";
 
+  const bool group_dispatch = op->GetAttr<bool>(ir::kGroupNextLevelDispatchAttr, false);
+  INTERNAL_CHECK_SPAN(!group_dispatch || !group_dispatch_active_, op->span_)
+      << "Nested grouped next-level dispatch loops are not supported";
+  if (group_dispatch) {
+    const std::string suffix = std::to_string(group_dispatch_counter_++);
+    group_dispatch_args_var_ = "_group_args_" + suffix;
+    group_dispatch_workers_var_ = "_group_workers_" + suffix;
+    group_dispatch_callee_.clear();
+    emitter_.EmitLine(group_dispatch_args_var_ + " = []");
+    emitter_.EmitLine(group_dispatch_workers_var_ + " = []");
+    group_dispatch_active_ = true;
+  }
+
   emitter_.EmitLine("for " + loop_var + " in range(" + start + ", " + stop + ", " + step + "):");
   emitter_.IncreaseIndent();
 
@@ -847,6 +866,16 @@ void DistributedCodegen::VisitStmt_(const ir::ForStmtPtr& op) {
   }
 
   emitter_.DecreaseIndent();
+  if (group_dispatch) {
+    group_dispatch_active_ = false;
+    INTERNAL_CHECK_SPAN(!group_dispatch_callee_.empty(), op->span_)
+        << "Grouped next-level dispatch loop emitted no CHIP-orchestrator call";
+    emitter_.EmitLine("_submit_chip_group(orch, callables[\"" + group_dispatch_callee_ + "\"], " +
+                      group_dispatch_args_var_ + ", config, " + group_dispatch_workers_var_ + ")");
+    group_dispatch_args_var_.clear();
+    group_dispatch_workers_var_.clear();
+    group_dispatch_callee_.clear();
+  }
 }
 
 void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) {
@@ -1205,8 +1234,19 @@ void DistributedCodegen::EmitCallToWorker(const ir::CallPtr& call, const ir::Fun
     // namespacing — see its docstring.
     emitter_.EmitLine("_keep.append(" + ta_var + ")");
     const std::string worker_arg = rank_expr.empty() ? "None" : rank_expr;
-    emitter_.EmitLine("_submit_chip(orch, callables[\"" + callee->name_ + "\"], " + ta_var + ", config, " +
-                      worker_arg + ")");
+    if (group_dispatch_active_) {
+      INTERNAL_CHECK_SPAN(!rank_expr.empty(), call->span_)
+          << "Grouped CHIP dispatch must carry an exact device rank";
+      INTERNAL_CHECK_SPAN(group_dispatch_callee_.empty() || group_dispatch_callee_ == callee->name_,
+                          call->span_)
+          << "Grouped dispatch loop must target exactly one CHIP orchestrator";
+      group_dispatch_callee_ = callee->name_;
+      emitter_.EmitLine(group_dispatch_args_var_ + ".append(" + ta_var + ")");
+      emitter_.EmitLine(group_dispatch_workers_var_ + ".append(" + worker_arg + ")");
+    } else {
+      emitter_.EmitLine("_submit_chip(orch, callables[\"" + callee->name_ + "\"], " + ta_var + ", config, " +
+                        worker_arg + ")");
+    }
   }
 
   // If this call has an assignment target (return value), alias it to the OUT

src/codegen/orchestration/orchestration_codegen.cpp

@@ -143,7 +143,11 @@ std::string GenerateIncludes(bool include_optional, bool include_vector = false)
     oss << "#include <optional>\n";
   }
   oss << "\n";
-  oss << "#include \"pto_orchestration_api.h\"\n\n";
+  oss << "#if __has_include(\"orchestration_api.h\")\n";
+  oss << "#include \"orchestration_api.h\"\n";
+  oss << "#else\n";
+  oss << "#include \"pto_orchestration_api.h\"\n";
+  oss << "#endif\n\n";
   return oss.str();
 }