diff --git a/docs/en/dev/codegen/01-orchestration_codegen.md b/docs/en/dev/codegen/01-orchestration_codegen.md index 7c100b0495..8a08b74216 100644 --- a/docs/en/dev/codegen/01-orchestration_codegen.md +++ b/docs/en/dev/codegen/01-orchestration_codegen.md @@ -73,7 +73,11 @@ This allows extensible operation codegen without modifying the core visitor. #include #include #include +#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 #include #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif extern "C" { diff --git a/docs/en/dev/passes/41-materialize_comm_domain_scopes.md b/docs/en/dev/passes/41-materialize_comm_domain_scopes.md index 37bad20a7b..270fef10cc 100644 --- a/docs/en/dev/passes/41-materialize_comm_domain_scopes.md +++ b/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"` so codegen emits the matching `__comm_d` 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 diff --git a/docs/zh/dev/codegen/01-orchestration_codegen.md b/docs/zh/dev/codegen/01-orchestration_codegen.md index 239a63b550..73cee8c45d 100644 --- a/docs/zh/dev/codegen/01-orchestration_codegen.md +++ b/docs/zh/dev/codegen/01-orchestration_codegen.md @@ -73,7 +73,11 @@ REGISTER_ORCHESTRATION_OP("tensor.slice", TensorSliceHandler); #include #include #include +#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 #include #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif extern "C" { diff --git a/docs/zh/dev/passes/41-materialize_comm_domain_scopes.md b/docs/zh/dev/passes/41-materialize_comm_domain_scopes.md index 9e800a4dcb..aa05ffafca 100644 --- a/docs/zh/dev/passes/41-materialize_comm_domain_scopes.md +++ b/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"`,使 codegen 能直接生成对应的 `__comm_d` 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 属性 diff --git a/include/pypto/codegen/distributed/distributed_codegen.h b/include/pypto/codegen/distributed/distributed_codegen.h index 02ebefec7e..aac079d535 100644 --- a/include/pypto/codegen/distributed/distributed_codegen.h +++ b/include/pypto/codegen/distributed/distributed_codegen.h @@ -299,6 +299,11 @@ class DistributedCodegen : public CodegenBase { std::set 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 diff --git a/include/pypto/ir/transforms/passes.h b/include/pypto/ir/transforms/passes.h index 94087a5fe3..fb4809af91 100644 --- a/include/pypto/ir/transforms/passes.h +++ b/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 diff --git a/include/pypto/ir/transforms/utils/attrs.h b/include/pypto/ir/transforms/utils/attrs.h index 5bddbc969a..c8636cae93 100644 --- a/include/pypto/ir/transforms/utils/attrs.h +++ b/include/pypto/ir/transforms/utils/attrs.h @@ -201,6 +201,13 @@ inline std::vector> 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``) // --------------------------------------------------------------------------- diff --git a/python/pypto/runtime/builtins/collectives/all_to_all/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/all_to_all/templates/entry.cpp.in index 31a78da027..fd2a0a52d2 100644 --- a/python/pypto/runtime/builtins/collectives/all_to_all/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/all_to_all/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in index 10b22f4f04..eda1f4d421 100644 --- a/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/allgather/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/allgather/templates/entry.cpp.in index 5316bb5371..b2a2c9c704 100644 --- a/python/pypto/runtime/builtins/collectives/allgather/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allgather/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in index 2346add36d..614a5f4fb6 100644 --- a/python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/entry.cpp.in index 2e9b19d506..229273d262 100644 --- a/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/barrier/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/barrier/templates/entry.cpp.in index 17a3299612..574cd193b7 100644 --- a/python/pypto/runtime/builtins/collectives/barrier/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/barrier/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/broadcast/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/broadcast/templates/entry.cpp.in index 76b5e2a68e..eaef207d72 100644 --- a/python/pypto/runtime/builtins/collectives/broadcast/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/broadcast/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/entry.cpp.in b/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/entry.cpp.in index a833afd2d4..f776ad501a 100644 --- a/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/entry.cpp.in +++ b/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/entry.cpp.in @@ -13,7 +13,11 @@ #include +#if __has_include("orchestration_api.h") +#include "orchestration_api.h" +#else #include "pto_orchestration_api.h" +#endif namespace { diff --git a/python/pypto/runtime/distributed_runner.py b/python/pypto/runtime/distributed_runner.py index 36bec5953e..c6d0ab0698 100644 --- a/python/pypto/runtime/distributed_runner.py +++ b/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. diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index 092777cc4d..b7afb391ec 100644 --- a/src/codegen/distributed/distributed_codegen.cpp +++ b/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`` (``/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(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 diff --git a/src/codegen/orchestration/orchestration_codegen.cpp b/src/codegen/orchestration/orchestration_codegen.cpp index 1301fec4ba..cbeaa9e94b 100644 --- a/src/codegen/orchestration/orchestration_codegen.cpp +++ b/src/codegen/orchestration/orchestration_codegen.cpp @@ -143,7 +143,11 @@ std::string GenerateIncludes(bool include_optional, bool include_vector = false) oss << "#include \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(); } diff --git a/src/ir/transforms/materialize_comm_domain_scopes_pass.cpp b/src/ir/transforms/materialize_comm_domain_scopes_pass.cpp index cbb529b935..b7e5b72c39 100644 --- a/src/ir/transforms/materialize_comm_domain_scopes_pass.cpp +++ b/src/ir/transforms/materialize_comm_domain_scopes_pass.cpp @@ -32,9 +32,11 @@ #include "pypto/ir/scalar_expr.h" #include "pypto/ir/span.h" #include "pypto/ir/stmt.h" +#include "pypto/ir/transforms/base/mutator.h" #include "pypto/ir/transforms/base/visitor.h" #include "pypto/ir/transforms/pass_properties.h" #include "pypto/ir/transforms/passes.h" +#include "pypto/ir/transforms/utils/attrs.h" #include "pypto/ir/transforms/utils/mutable_copy.h" #include "pypto/ir/transforms/utils/transform_utils.h" #include "pypto/ir/type.h" @@ -471,6 +473,89 @@ class DispatchAnalyzer : public IRVisitor { int repeating_scope_depth_ = 0; }; +/// Collect side effects and calls in one loop iteration without looking through +/// nested control flow. Group publication may reorder the CHIP dispatch after +/// argument construction, so only pure tensor view construction is allowed +/// beside one unconditional CHIP dispatch; any Submit is rejected explicitly. +class GroupDispatchBodyAnalyzer : public IRVisitor { + public: + explicit GroupDispatchBodyAnalyzer(const std::map& chip_orchs) + : chip_orchs_(chip_orchs) {} + + void VisitStmt_(const ForStmtPtr& /*op*/) override { has_nested_control = true; } + void VisitStmt_(const WhileStmtPtr& /*op*/) override { has_nested_control = true; } + void VisitStmt_(const IfStmtPtr& /*op*/) override { has_nested_control = true; } + void VisitExpr_(const SubmitPtr& /*op*/) override { has_submit = true; } + + void VisitExpr_(const CallPtr& op) override { + if (IsChipOrchDispatch(op, chip_orchs_)) { + dispatches.push_back(op); + } else if (!IsOp(op, "tensor.slice") && !IsOp(op, "pld.tensor.window")) { + has_other_call = true; + } + IRVisitor::VisitExpr_(op); + } + + bool has_nested_control{false}; + bool has_submit{false}; + bool has_other_call{false}; + std::vector dispatches; + + private: + const std::map& chip_orchs_; +}; + +/// Mark only compiler-proven full-world rank loops for group publication. +/// Codegen deliberately does not rediscover this pattern: the pass owns the +/// semantic proof, and the loop attr is the explicit lowering contract. +class GroupDispatchLoopMarker : public IRMutator { + public: + GroupDispatchLoopMarker(const std::map& chip_orchs, + const std::unordered_map& var_defs) + : chip_orchs_(chip_orchs), var_defs_(var_defs) {} + + protected: + StmtPtr VisitStmt_(const ForStmtPtr& op) override { + auto rewritten = As(IRMutator::VisitStmt_(op)); + INTERNAL_CHECK(rewritten); + if (!IsEligible(rewritten) || rewritten->GetAttr(kGroupNextLevelDispatchAttr, false)) { + return rewritten; + } + auto marked = MutableCopy(rewritten); + marked->attrs_.emplace_back(kGroupNextLevelDispatchAttr, true); + return marked; + } + + private: + [[nodiscard]] bool IsEligible(const ForStmtPtr& op) const { + if (!op->iter_args_.empty() || !op->return_vars_.empty()) return false; + auto start = As(UnwrapStopExpr(op->start_, var_defs_)); + auto step = As(UnwrapStopExpr(op->step_, var_defs_)); + if (!start || start->value_ != 0 || !step || step->value_ != 1) return false; + auto stop = As(UnwrapStopExpr(op->stop_, var_defs_)); + if (!stop || !stop->op_ || !IsOp(stop, "pld.system.world_size")) return false; + + GroupDispatchBodyAnalyzer analyzer(chip_orchs_); + analyzer.VisitStmt(op->body_); + if (analyzer.has_nested_control || analyzer.has_submit || analyzer.has_other_call || + analyzer.dispatches.size() != 1) { + return false; + } + ExprPtr device; + for (const auto& [key, value] : analyzer.dispatches.front()->attrs_) { + if (key == kAttrDevice) { + if (const auto* expr = std::any_cast(&value)) device = *expr; + break; + } + } + auto device_var = As(device); + return device_var && device_var.get() == op->loop_var_.get(); + } + + const std::map& chip_orchs_; + const std::unordered_map& var_defs_; +}; + /// A host-orchestration function in PyPTO is declared as either /// ``@pl.function(type=FunctionType.Orchestration, level=Level.HOST)`` or /// (more common in distributed programs) ``@pl.function(level=Level.HOST, @@ -528,9 +613,17 @@ FunctionPtr ProcessHostOrch(const FunctionPtr& func, const std::mapbody_.get()) return func; + auto marked_func = MutableCopy(func); + marked_func->body_ = materialization_body; + return marked_func; } // Phase 2: record device-descriptor evidence from dispatch sites. diff --git a/tests/ut/codegen/distributed/test_host_orch_distributed.py b/tests/ut/codegen/distributed/test_host_orch_distributed.py index ef7969cdae..8d0720af52 100644 --- a/tests/ut/codegen/distributed/test_host_orch_distributed.py +++ b/tests/ut/codegen/distributed/test_host_orch_distributed.py @@ -19,8 +19,9 @@ 3. Explicit CommCtx scalar: ``add_scalar(__comm_d0[].device_ctx)`` placed AFTER all tensor adds, in IR-arg order (matching the materialized incore function signature). -4. dispatch ``device=`` attr → ``_submit_chip(orch, ..., config, )`` (the - rank-pinned wrapper that namespaces per-rank DFX ``output_prefix``). +4. dispatch ``device=`` attr → a rank-pinned submission; compiler-proven + full-world loops use ``_submit_chip_group``, while other shapes use + ``_submit_chip``. Plus regressions: @@ -229,9 +230,150 @@ def host_orch( # the Buffer.tensor above. assert re.search(r"\.add_scalar\(__comm_d0\[\w+\]\.device_ctx\)", code), code assert "pld.system.get_comm_ctx" not in code, code - # ``device=r`` → rank-pinned dispatch routes through ``_submit_chip`` (which - # namespaces the per-rank DFX ``output_prefix``), passing the rank last. - assert re.search(r"_submit_chip\(orch, callables\[\"chip_orch\"\],.*config, \w+\)", code), code + # A full-world one-dispatch loop is published as one group after all ranks' + # TaskArgs have been built. + assert re.search(r"_group_workers_0\.append\(\w+\)", code), code + assert '_submit_chip_group(orch, callables["chip_orch"]' in code, code + + +def test_full_rank_dispatch_loop_batches_task_args_before_submit(): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def chip_orch( + self, + data: pld.DistributedTensor[[SIZE], pl.FP32], + ) -> pl.Tensor[[SIZE], pl.FP32]: + return data # type: ignore[return-value] + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch(self) -> pl.Tensor[[SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + data = pld.window(data_buf, [SIZE], dtype=pl.FP32) + for r in pl.range(pld.world_size()): + self.chip_orch(data, device=r) + return data # type: ignore[return-value] + + code = _lower(Prog) + + loop_end = code.index("_submit_chip_group(") + assert "_group_args_0.append(_ta_0)" in code[:loop_end], code + assert "_group_workers_0.append(" in code[:loop_end], code + assert re.search( + r'_submit_chip_group\(orch, callables\["chip_orch"\], _group_args_0, ' + r"config, _group_workers_0\)", + code[loop_end:], + ), code + assert '_submit_chip(orch, callables["chip_orch"], _ta_0' not in code, code + compile(code, "", "exec") + + +def test_allocation_free_full_rank_dispatch_loop_is_grouped(): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def chip_orch( + self, + inp: pl.Tensor[[SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[SIZE], pl.FP32]], + ) -> pl.Tensor[[SIZE], pl.FP32]: + return out + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[2, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[2, SIZE], pl.FP32]], + ) -> pl.Tensor[[2, SIZE], pl.FP32]: + for r in pl.range(pld.world_size()): + inp_r = inputs[r] + out_r = outputs[r] + self.chip_orch(inp_r, out_r, device=r) + return outputs + + code = _lower(Prog) + + assert "allocate_domain" not in code, code + assert '_submit_chip_group(orch, callables["chip_orch"]' in code, code + assert '_submit_chip(orch, callables["chip_orch"]' not in code, code + compile(code, "", "exec") + + +def test_full_rank_dispatch_loop_allows_per_rank_window_views(): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def chip_orch( + self, + data: pld.DistributedTensor[[SIZE], pl.FP32], + ) -> pl.Tensor[[SIZE], pl.FP32]: + return data # type: ignore[return-value] + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch(self) -> pl.Tensor[[SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [SIZE], dtype=pl.FP32) + self.chip_orch(data, device=r) + return data # type: ignore[return-value] + + code = _lower(Prog) + + assert '_submit_chip_group(orch, callables["chip_orch"]' in code, code + assert '_submit_chip(orch, callables["chip_orch"]' not in code, code + compile(code, "", "exec") + + +def test_multiple_dispatches_in_rank_loop_remain_individual_submissions(): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def chip_orch( + self, + data: pld.DistributedTensor[[SIZE], pl.FP32], + ) -> pl.Tensor[[SIZE], pl.FP32]: + return data # type: ignore[return-value] + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch(self) -> pl.Tensor[[SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + data = pld.window(data_buf, [SIZE], dtype=pl.FP32) + for r in pl.range(pld.world_size()): + self.chip_orch(data, device=r) + self.chip_orch(data, device=r) + return data # type: ignore[return-value] + + code = _lower(Prog) + + assert "_submit_chip_group(" not in code, code + assert code.count('_submit_chip(orch, callables["chip_orch"]') == 2, code + compile(code, "", "exec") + + +def test_other_call_in_rank_loop_keeps_program_order(): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.Orchestration) + def chip_orch( + self, + data: pld.DistributedTensor[[SIZE], pl.FP32], + ) -> pl.Tensor[[SIZE], pl.FP32]: + return data # type: ignore[return-value] + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch(self) -> pl.Tensor[[SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + data = pld.window(data_buf, [SIZE], dtype=pl.FP32) + for r in pl.range(pld.world_size()): + _world_size = pld.world_size() + self.chip_orch(data, device=r) + return data # type: ignore[return-value] + + code = _lower(Prog) + + assert "_submit_chip_group(" not in code, code + assert code.count('_submit_chip(orch, callables["chip_orch"]') == 1, code + compile(code, "", "exec") def test_two_dist_tensor_formals_emit_two_explicit_ctx_scalars(): diff --git a/tests/ut/codegen/test_orchestration_codegen.py b/tests/ut/codegen/test_orchestration_codegen.py index d33655ec38..3a94d95aa3 100644 --- a/tests/ut/codegen/test_orchestration_codegen.py +++ b/tests/ut/codegen/test_orchestration_codegen.py @@ -68,7 +68,11 @@ def orch_basic( #include #include + #if __has_include("orchestration_api.h") + #include "orchestration_api.h" + #else #include "pto_orchestration_api.h" + #endif extern "C" { @@ -517,7 +521,11 @@ def orch_vector( #include #include + #if __has_include("orchestration_api.h") + #include "orchestration_api.h" + #else #include "pto_orchestration_api.h" + #endif extern "C" { @@ -1555,7 +1563,11 @@ def orch_inplace( #include #include + #if __has_include("orchestration_api.h") + #include "orchestration_api.h" + #else #include "pto_orchestration_api.h" + #endif extern "C" { diff --git a/tests/ut/runtime/test_distributed_worker.py b/tests/ut/runtime/test_distributed_worker.py index b5b765c716..965f12793b 100644 --- a/tests/ut/runtime/test_distributed_worker.py +++ b/tests/ut/runtime/test_distributed_worker.py @@ -53,6 +53,7 @@ _make_call_config, _reset_dfx_dispatch_state, _submit_chip, + _submit_chip_group, ) from pypto.runtime.runner import RunConfig @@ -2075,6 +2076,7 @@ class _RecordingOrch: def __init__(self, chip_count: int | None = None) -> None: self.calls: list[tuple[Any, int, str]] = [] + self.group_call: tuple[Any, list[Any], str, list[int]] | None = None # ``_submit_chip`` reads/writes this per-card dispatch counter on the # orch; declare it so the attribute is known to the type checker. self._dfx_dispatch_idx: dict[str, int] = {} @@ -2089,6 +2091,12 @@ def submit_next_level(self, callable_id: Any, task_args: Any, config: Any, *, wo self.calls.append((callable_id, worker, config.output_prefix)) return "submitted" + def submit_next_level_group( + self, callable_id: Any, task_args_list: list[Any], config: Any, *, workers: list[int] + ) -> str: + self.group_call = (callable_id, task_args_list, config.output_prefix, workers) + return "group-submitted" + class TestSubmitChip: """``_submit_chip`` namespaces per-dispatch DFX ``output_prefix`` then restores it.""" @@ -2224,6 +2232,41 @@ def test_no_marker_when_chip_names_unstamped(self, tmp_path): assert not (tmp_path / "rank0" / "d0" / "dispatch_program.json").exists() +class TestSubmitChipGroup: + """Grouped dispatch synchronizes publication without breaking DFX isolation.""" + + def test_dfx_off_uses_one_group_submission(self): + orch = _RecordingOrch() + cfg = _SpyDfxConfig(output_prefix="") + + ret = _submit_chip_group(orch, "chip", ["ta0", "ta1"], cfg, [0, 1]) + + assert ret == "group-submitted" + assert orch.group_call == ("chip", ["ta0", "ta1"], "", [0, 1]) + assert orch.calls == [] + + def test_dfx_on_preserves_per_member_namespaces(self): + orch = _RecordingOrch() + cfg = _SpyDfxConfig(output_prefix="/work/dfx_outputs") + + ret = _submit_chip_group(orch, "chip", ["ta0", "ta1"], cfg, [0, 1]) + + assert ret == ["submitted", "submitted"] + assert orch.calls == [ + ("chip", 0, "/work/dfx_outputs/rank0/d0"), + ("chip", 1, "/work/dfx_outputs/rank1/d0"), + ] + assert cfg.output_prefix == "/work/dfx_outputs" + + @pytest.mark.parametrize( + ("task_args_list", "workers", "message"), + [([], [], "at least one"), (["ta0"], [0, 1], "length"), (["ta0", "ta1"], [0, 0], "duplicate")], + ) + def test_rejects_invalid_groups(self, task_args_list, workers, message): + with pytest.raises(ValueError, match=message): + _submit_chip_group(_RecordingOrch(), "chip", task_args_list, _SpyDfxConfig(), workers) + + def _write_dfx_dispatch_dirs(dfx: Path, *rels: str) -> None: """Lay down ``//chip_swimlane_records.json`` for each dispatch dir.