From ecbf7e1acd45070844ccacadae3ed20ee0952a22 Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Tue, 28 Jul 2026 12:57:26 +0200 Subject: [PATCH 1/4] fix(codegen): emit tensor phis via tensors dict to avoid NameError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When ConvertToSSA synthesizes phi `return_vars_` for cross-branch diverging tensor variables, host_orch.py codegen must pre-declare the phi name in the `tensors` dict and emit yield-to-phi assignments that reference tensors-dict names — not bare Python variables. Fixes #2180. - IfStmt visitor (DistributedCodegen): pre-declares phi variables before `if`, emits branch-specific yield-to-phi `tensors[…] = tensors[…]` assignments, and traces Var→Var aliases (kernel param copies) to the original tensors-dict name. - AssignStmt visitor: aliases `tensor_var = kernel_param` are emitted as `tensors["tensor_var"] = tensors["kernel_param"]` instead of bare Python names. - Unit test validates phi pre-declaration, yield assignments, and absence of bare-name tensor assignments. Co-authored-by: georgebisbas Co-authored-by: vloncar --- .../distributed/distributed_codegen.cpp | 146 ++++++++++++++++++ .../distributed/test_host_orch_distributed.py | 71 ++++++++- 2 files changed, 216 insertions(+), 1 deletion(-) diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index f418204129..b745da1e42 100644 --- a/src/codegen/distributed/distributed_codegen.cpp +++ b/src/codegen/distributed/distributed_codegen.cpp @@ -38,6 +38,7 @@ #include "pypto/ir/program.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/stmt.h" +#include "pypto/ir/transforms/utils/transform_utils.h" #include "pypto/ir/type.h" namespace pypto { @@ -743,6 +744,21 @@ void DistributedCodegen::VisitStmt_(const ir::AssignStmtPtr& op) { return; } + // Tensor→tensor alias: ``t = kernel_param`` — emit as a tensors-dict + // reference so both names resolve via ``tensors[...]`` in the generated + // Python. Bare Python names for tensors are invisible to the prepared + // runtime's tensor registry (issue #2180). + if (ir::AsTensorTypeLike(op->var_->GetType()) && ir::As(op->value_)) { + VisitExpr(op->value_); + if (!current_expr_value_.empty()) { + emitter_.EmitLine("tensors[\"" + var_name + "\"] = tensors[\"" + current_expr_value_ + "\"]"); + declared_vars_.insert(var_name); + current_expr_value_ = ""; + } + current_target_var_ = ""; + return; + } + // Standard expression VisitExpr(op->value_); @@ -813,6 +829,136 @@ void DistributedCodegen::VisitStmt_(const ir::ForStmtPtr& op) { void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { INTERNAL_CHECK(op != nullptr) << "Internal error: null IfStmt"; + // When ConvertToSSA synthesized phi return_vars_ for diverging variables + // across branches, we must emit pre-declarations AND yield-to-phi + // assignments so the phi name is visible to post-if consumers. Without + // this a branch-local SSA name (defined inside one branch's Python scope) + // leaks into the merged scope and raises NameError at prepare(). + if (!op->return_vars_.empty()) { + // ── Extract trailing yields from both branches ────────────────── + const auto then_yield = ir::transform_utils::GetLastYieldStmt( + ir::transform_utils::UnwrapAutoScope(op->then_body_)); + const ir::YieldStmtPtr else_yield = [&]() -> ir::YieldStmtPtr { + if (!op->else_body_.has_value()) return nullptr; + return ir::transform_utils::GetLastYieldStmt( + ir::transform_utils::UnwrapAutoScope(*op->else_body_)); + }(); + + // ── Helper: find the AssignStmt that defines a yield var in a branch ── + // The yield value is the SSA name of the diverging variable inside the + // branch. For tensor-typed phis, we need to resolve this to the + // tensor's actual source — the RHS of the assignment that sets it. + // This matters for the then-branch where ``boundary = zero`` creates a + // bare Python name that is NOT in the ``tensors`` dict, while + // ``zero`` (a kernel param) IS. + auto find_yield_source = [](const ir::StmtPtr& branch_body, + const ir::VarPtr& yield_var) -> ir::ExprPtr { + const auto stmts = ir::transform_utils::FlattenToStmts( + ir::transform_utils::UnwrapAutoScope(branch_body)); + // Walk backwards (excluding the trailing YieldStmt) to find the + // most recent assignment to this var. + for (auto it = stmts.rbegin(); it != stmts.rend(); ++it) { + if (ir::As(*it)) continue; + if (auto assign = ir::As(*it)) { + if (assign->var_ == yield_var) return assign->value_; + } + } + return nullptr; + }; + + // ── Find a tensor init for tensor-typed phis ──────────────────── + // Tensor phis need an in-scope init to reference in the tensors dict. + // Scan function params first (guaranteed in tensors dict at runtime), + // then fall back to the then-branch yield's first tensor value. + std::string tensor_phi_init; + for (const auto& param : current_func_->params_) { + if (ir::AsTensorTypeLike(param->GetType())) { + tensor_phi_init = SanitizeName(param->name_hint_); + break; + } + } + if (tensor_phi_init.empty() && then_yield) { + for (const auto& val : then_yield->value_) { + const auto var = ir::As(val); + if (var && ir::AsTensorTypeLike(var->GetType())) { + auto src = find_yield_source(op->then_body_, var); + if (src) { + VisitExpr(src); + tensor_phi_init = current_expr_value_; + current_expr_value_ = ""; + } + break; + } + } + } + + // ── Pre-declare phi variables before the ``if`` ───────────────── + for (size_t i = 0; i < op->return_vars_.size(); ++i) { + const std::string phi_name = SanitizeName(op->return_vars_[i]->name_hint_); + if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) { + if (!tensor_phi_init.empty()) { + emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + tensor_phi_init + "\"]"); + } else { + emitter_.EmitLine("if \"" + phi_name + "\" not in tensors:"); + emitter_.EmitLine(" tensors[\"" + phi_name + "\"] = torch.zeros((1,), dtype=torch.float32).share_memory_()"); + } + } else { + emitter_.EmitLine(phi_name + " = None"); + } + declared_vars_.insert(phi_name); + } + + // ── Emit below-branch yield-to-phi assignments ───────────────── + // Use the assignment RHS as the yield source when the RHS is a Var + // (e.g. ``boundary = zero`` — a kernel-param alias). For Call RHS + // (e.g. ``boundary = self.chip_run(...)``), EmitCallToWorker already + // places the result in ``tensors[yield_var_name]``, so use the yield + // var name directly. + auto emit_branch_yields = [&](const ir::StmtPtr& branch_body, const ir::YieldStmtPtr& yld) { + if (!yld) return; + for (size_t i = 0; i < op->return_vars_.size() && i < yld->value_.size(); ++i) { + const auto yield_var = ir::As(yld->value_[i]); + if (!yield_var) continue; + + // Find the assignment that defines this yield var. + auto src = find_yield_source(branch_body, yield_var); + // Only trace Var→Var aliases. For Call RHS the result is already + // in ``tensors[yield_var_name]`` via EmitCallToWorker. + const bool is_var_alias = src && ir::As(src) != nullptr; + const ir::ExprPtr yield_expr = is_var_alias ? src : yield_var; + VisitExpr(yield_expr); + const std::string yield_val = current_expr_value_; + current_expr_value_ = ""; + const std::string phi_name = SanitizeName(op->return_vars_[i]->name_hint_); + if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) { + emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + yield_val + "\"]"); + } else { + emitter_.EmitLine(phi_name + " = " + yield_val); + } + } + }; + // ── Emit condition and if/else with yield-to-phi ──────────────── + VisitExpr(op->condition_); + std::string condition = current_expr_value_; + current_expr_value_ = ""; + + emitter_.EmitLine("if " + condition + ":"); + emitter_.IncreaseIndent(); + VisitStmt(op->then_body_); + emit_branch_yields(op->then_body_, then_yield); + emitter_.DecreaseIndent(); + + if (op->else_body_.has_value()) { + emitter_.EmitLine("else:"); + emitter_.IncreaseIndent(); + VisitStmt(*op->else_body_); + emit_branch_yields(*op->else_body_, else_yield); + emitter_.DecreaseIndent(); + } + return; + } + + // ── Simple if/else without return_vars_ (no phi nodes) ─────────── VisitExpr(op->condition_); std::string condition = current_expr_value_; current_expr_value_ = ""; diff --git a/tests/ut/codegen/distributed/test_host_orch_distributed.py b/tests/ut/codegen/distributed/test_host_orch_distributed.py index 9219425283..a631241351 100644 --- a/tests/ut/codegen/distributed/test_host_orch_distributed.py +++ b/tests/ut/codegen/distributed/test_host_orch_distributed.py @@ -67,8 +67,10 @@ def pass_verification_context(): yield -def _lower(program) -> str: +def _lower(program, convert_to_ssa=False) -> str: """Apply the late host-distributed pipeline, then run distributed codegen directly.""" + if convert_to_ssa: + program = passes.convert_to_ssa()(program) program = passes.synthesize_allreduce_signals()(program) program = passes.materialize_comm_domain_scopes()(program) program = passes.materialize_dist_tensor_ctx()(program) @@ -949,5 +951,72 @@ def test_host_collective_builtin_template_package_exists(package_name, variant): assert variant.startswith("builtin.tensor."), variant +# --------------------------------------------------------------------------- +# IfStmt phi codegen (issue #2180) +# --------------------------------------------------------------------------- + + +def test_if_cross_branch_phi_predeclares_and_yields_tensors() -> None: + """Tensor phi emitted by ConvertToSSA for a cross-branch diverging variable + is pre-declared and yielded via ``tensors[...]``, not bare Python names.""" + SIZE = 4 + P = 2 + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def identity(self, x: pl.Tensor[[SIZE, SIZE], pl.FP32]) -> pl.Tensor[[SIZE, SIZE], pl.FP32]: + return x + + @pl.function(type=pl.FunctionType.Orchestration) + def chip_run(self, x: pl.Tensor[[SIZE, SIZE], pl.FP32]) -> pl.Tensor[[SIZE, SIZE], pl.FP32]: + return self.identity(x) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch(self, x: pl.Tensor[[SIZE, SIZE], pl.FP32], + zero: pl.Tensor[[SIZE, SIZE], pl.FP32]) -> pl.Scalar[pl.INT32]: + for r in pl.range(P): + if r == 0: + boundary = zero + else: + boundary = self.chip_run(x) + self.chip_run(boundary) + return 0 + + code = _lower(Prog, convert_to_ssa=True) + + # Pre-declared phi before the ``if`` — must use tensors[...] for tensor phis. + assert re.search(r'tensors\["boundary__phi_v\d+"\]\s*=.*tensors', code), ( + "Tensor phi must be pre-declared via tensors[...] before if block:\n" + code + ) + + # Yield in the then-branch must reference a valid tensors-dict entry (no + # bare Python names for tensor phis). + then_yield = re.search( + r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', + re.split(r"\belse\b:", code)[0], + ) + assert then_yield is not None, ( + "Then-branch must yield phi via tensors[...] = tensors[...]:\n" + code + ) + + # Yield in the else-branch likewise. + else_block = code.split("else:")[1] if "else:" in code else "" + else_yield = re.search( + r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', + else_block, + ) + assert else_yield is not None, ( + "Else-branch must yield phi via tensors[...] = tensors[...]:\n" + code + ) + + # No bare-Python-name tensor assignment in the then-branch (the bug was + # ``boundary__ssa_v0 = zero__ssa_v0`` which raised NameError). + then_block = re.split(r"\belse\b:", code)[0] + assert not re.search(r"(? Date: Tue, 28 Jul 2026 13:37:55 +0200 Subject: [PATCH 2/4] fix(codegen): per-phi init, test assertion, and pre-commit (#2183) CodeRabbit review: each tensor-typed phi now derives its pre-declaration initializer from its own yield source instead of a shared init. Test assertion isolates the then-branch body from the pre-if declaration. Pre-commit: suppress pyright reportReturnType for IR-level return 0. --- .../distributed/distributed_codegen.cpp | 36 +++++++++++++------ .../distributed/test_host_orch_distributed.py | 33 ++++++++++------- 2 files changed, 46 insertions(+), 23 deletions(-) diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index b745da1e42..4a22686ced 100644 --- a/src/codegen/distributed/distributed_codegen.cpp +++ b/src/codegen/distributed/distributed_codegen.cpp @@ -836,12 +836,11 @@ void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { // leaks into the merged scope and raises NameError at prepare(). if (!op->return_vars_.empty()) { // ── Extract trailing yields from both branches ────────────────── - const auto then_yield = ir::transform_utils::GetLastYieldStmt( - ir::transform_utils::UnwrapAutoScope(op->then_body_)); + const auto then_yield = + ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(op->then_body_)); const ir::YieldStmtPtr else_yield = [&]() -> ir::YieldStmtPtr { if (!op->else_body_.has_value()) return nullptr; - return ir::transform_utils::GetLastYieldStmt( - ir::transform_utils::UnwrapAutoScope(*op->else_body_)); + return ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(*op->else_body_)); }(); // ── Helper: find the AssignStmt that defines a yield var in a branch ── @@ -851,10 +850,9 @@ void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { // This matters for the then-branch where ``boundary = zero`` creates a // bare Python name that is NOT in the ``tensors`` dict, while // ``zero`` (a kernel param) IS. - auto find_yield_source = [](const ir::StmtPtr& branch_body, - const ir::VarPtr& yield_var) -> ir::ExprPtr { - const auto stmts = ir::transform_utils::FlattenToStmts( - ir::transform_utils::UnwrapAutoScope(branch_body)); + auto find_yield_source = [](const ir::StmtPtr& branch_body, const ir::VarPtr& yield_var) -> ir::ExprPtr { + const auto stmts = + ir::transform_utils::FlattenToStmts(ir::transform_utils::UnwrapAutoScope(branch_body)); // Walk backwards (excluding the trailing YieldStmt) to find the // most recent assignment to this var. for (auto it = stmts.rbegin(); it != stmts.rend(); ++it) { @@ -893,14 +891,32 @@ void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { } // ── Pre-declare phi variables before the ``if`` ───────────────── + // Derive each tensor phi's initializer from its own yield source in + // the then-branch (preserving shape/dtype), falling back to the + // shared param-derived init or a zero-placeholder. for (size_t i = 0; i < op->return_vars_.size(); ++i) { const std::string phi_name = SanitizeName(op->return_vars_[i]->name_hint_); if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) { - if (!tensor_phi_init.empty()) { + // Try a per-phi init from the then-branch yield first. + std::string per_phi_init; + if (then_yield && i < then_yield->value_.size()) { + const auto yield_var = ir::As(then_yield->value_[i]); + if (yield_var) { + auto src = find_yield_source(op->then_body_, yield_var); + const bool is_var_alias = src && ir::As(src) != nullptr; + VisitExpr(is_var_alias ? src : yield_var); + per_phi_init = current_expr_value_; + current_expr_value_ = ""; + } + } + if (!per_phi_init.empty()) { + emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + per_phi_init + "\"]"); + } else if (!tensor_phi_init.empty()) { emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + tensor_phi_init + "\"]"); } else { emitter_.EmitLine("if \"" + phi_name + "\" not in tensors:"); - emitter_.EmitLine(" tensors[\"" + phi_name + "\"] = torch.zeros((1,), dtype=torch.float32).share_memory_()"); + emitter_.EmitLine(" tensors[\"" + phi_name + + "\"] = torch.zeros((1,), dtype=torch.float32).share_memory_()"); } } else { emitter_.EmitLine(phi_name + " = None"); diff --git a/tests/ut/codegen/distributed/test_host_orch_distributed.py b/tests/ut/codegen/distributed/test_host_orch_distributed.py index a631241351..5fb1777271 100644 --- a/tests/ut/codegen/distributed/test_host_orch_distributed.py +++ b/tests/ut/codegen/distributed/test_host_orch_distributed.py @@ -973,46 +973,53 @@ def chip_run(self, x: pl.Tensor[[SIZE, SIZE], pl.FP32]) -> pl.Tensor[[SIZE, SIZE return self.identity(x) @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) - def host_orch(self, x: pl.Tensor[[SIZE, SIZE], pl.FP32], - zero: pl.Tensor[[SIZE, SIZE], pl.FP32]) -> pl.Scalar[pl.INT32]: + def host_orch( + self, x: pl.Tensor[[SIZE, SIZE], pl.FP32], zero: pl.Tensor[[SIZE, SIZE], pl.FP32] + ) -> pl.Scalar[pl.INT32]: for r in pl.range(P): if r == 0: boundary = zero else: boundary = self.chip_run(x) self.chip_run(boundary) - return 0 + return 0 # pyright: ignore[reportReturnType] code = _lower(Prog, convert_to_ssa=True) # Pre-declared phi before the ``if`` — must use tensors[...] for tensor phis. - assert re.search(r'tensors\["boundary__phi_v\d+"\]\s*=.*tensors', code), ( + pre_if = code.split("if", 1)[0] + assert re.search(r'tensors\["boundary__phi_v\d+"\]\s*=.*tensors', pre_if), ( "Tensor phi must be pre-declared via tensors[...] before if block:\n" + code ) + # Isolate the then-branch body: everything after the ``if ...:`` line header + # (first newline after the ``if`` keyword) up to the ``else:`` boundary. + if "else:" in code: + post_if = code.split("if", 1)[1].split("\n", 1)[1] # strip ``:\n`` header + then_block = post_if.split("else:")[0] + else_block = post_if.split("else:")[1] + else: + post_if = code.split("if", 1)[1].split("\n", 1)[1] + then_block = post_if + else_block = "" + # Yield in the then-branch must reference a valid tensors-dict entry (no # bare Python names for tensor phis). then_yield = re.search( r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', - re.split(r"\belse\b:", code)[0], - ) - assert then_yield is not None, ( - "Then-branch must yield phi via tensors[...] = tensors[...]:\n" + code + then_block, ) + assert then_yield is not None, "Then-branch must yield phi via tensors[...] = tensors[...]:\n" + code # Yield in the else-branch likewise. - else_block = code.split("else:")[1] if "else:" in code else "" else_yield = re.search( r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', else_block, ) - assert else_yield is not None, ( - "Else-branch must yield phi via tensors[...] = tensors[...]:\n" + code - ) + assert else_yield is not None, "Else-branch must yield phi via tensors[...] = tensors[...]:\n" + code # No bare-Python-name tensor assignment in the then-branch (the bug was # ``boundary__ssa_v0 = zero__ssa_v0`` which raised NameError). - then_block = re.split(r"\belse\b:", code)[0] assert not re.search(r"(? Date: Tue, 28 Jul 2026 14:33:12 +0200 Subject: [PATCH 3/4] fix(codegen): skip VisitExpr on Call RHS during phi init discovery (#2183) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a HOST orchestrator has no tensor formal parameters, the phi init fallback visited the then-branch yield source via VisitExpr. If that source is a hierarchy call, EmitCallToWorker would emit a _submit_chip before the `if` condition, altering program semantics. Only resolve Var→Var aliases now; Call RHS falls through to the zero placeholder. Co-authored-by: georgebisbas Co-authored-by: vloncar --- .../distributed/distributed_codegen.cpp | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index 4a22686ced..a21a7c19e0 100644 --- a/src/codegen/distributed/distributed_codegen.cpp +++ b/src/codegen/distributed/distributed_codegen.cpp @@ -880,7 +880,10 @@ void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { const auto var = ir::As(val); if (var && ir::AsTensorTypeLike(var->GetType())) { auto src = find_yield_source(op->then_body_, var); - if (src) { + // Only resolve Var→Var aliases for init discovery. Call RHS + // triggers EmitCallToWorker which would submit a chip task + // before the `if` condition is evaluated. + if (src && ir::As(src) != nullptr) { VisitExpr(src); tensor_phi_init = current_expr_value_; current_expr_value_ = ""; @@ -903,10 +906,15 @@ void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { const auto yield_var = ir::As(then_yield->value_[i]); if (yield_var) { auto src = find_yield_source(op->then_body_, yield_var); - const bool is_var_alias = src && ir::As(src) != nullptr; - VisitExpr(is_var_alias ? src : yield_var); - per_phi_init = current_expr_value_; - current_expr_value_ = ""; + // Only visit Var→Var aliases (e.g. ``boundary = zero``). + // For Call RHS, EmitCallToWorker would submit a chip task + // before the `if` condition — skip init discovery and fall + // through to the shared tensor_phi_init or zero placeholder. + if (src && ir::As(src) != nullptr) { + VisitExpr(src); + per_phi_init = current_expr_value_; + current_expr_value_ = ""; + } } } if (!per_phi_init.empty()) { From 040a4b0a65d1e06fc2cfa56a51be1df0e4da2cb1 Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Wed, 29 Jul 2026 10:07:15 +0000 Subject: [PATCH 4/4] fix(codegen): simplify IfStmt phi emission in distributed codegen Drop the pre-declaration loop, find_yield_source helper, and fallback initializers (torch.zeros placeholder and tensor_phi_init) from VisitStmt_(IfStmtPtr). The AssignStmt fix (also in this branch) guarantees that tensors[] is populated for every yield var by the time VisitStmt returns, so yields can be emitted directly with no RHS back-tracing. Also tighten the AssignStmt tensor-alias guard: use AsVarLike (vs As) to catch IterArg loop-carried tensors, and exact As (vs AsTensorTypeLike) to avoid over-matching DistributedTensorType. Test: replace fragile substring split with line-anchored if/else matching and add compile() sanity check. Rename to reflect simplified behavior. --- .../distributed/distributed_codegen.cpp | 181 +++--------------- .../distributed/test_host_orch_distributed.py | 68 +++---- 2 files changed, 64 insertions(+), 185 deletions(-) diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index a21a7c19e0..965875457f 100644 --- a/src/codegen/distributed/distributed_codegen.cpp +++ b/src/codegen/distributed/distributed_codegen.cpp @@ -748,7 +748,9 @@ void DistributedCodegen::VisitStmt_(const ir::AssignStmtPtr& op) { // reference so both names resolve via ``tensors[...]`` in the generated // Python. Bare Python names for tensors are invisible to the prepared // runtime's tensor registry (issue #2180). - if (ir::AsTensorTypeLike(op->var_->GetType()) && ir::As(op->value_)) { + // Only plain TensorType (not DistributedTensorType) goes through the + // tensors dict; distributed tensors use window_buffer_ references. + if (ir::As(op->var_->GetType()) && ir::AsVarLike(op->value_)) { VisitExpr(op->value_); if (!current_expr_value_.empty()) { emitter_.EmitLine("tensors[\"" + var_name + "\"] = tensors[\"" + current_expr_value_ + "\"]"); @@ -829,173 +831,46 @@ void DistributedCodegen::VisitStmt_(const ir::ForStmtPtr& op) { void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { INTERNAL_CHECK(op != nullptr) << "Internal error: null IfStmt"; - // When ConvertToSSA synthesized phi return_vars_ for diverging variables - // across branches, we must emit pre-declarations AND yield-to-phi - // assignments so the phi name is visible to post-if consumers. Without - // this a branch-local SSA name (defined inside one branch's Python scope) - // leaks into the merged scope and raises NameError at prepare(). - if (!op->return_vars_.empty()) { - // ── Extract trailing yields from both branches ────────────────── - const auto then_yield = - ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(op->then_body_)); - const ir::YieldStmtPtr else_yield = [&]() -> ir::YieldStmtPtr { - if (!op->else_body_.has_value()) return nullptr; - return ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(*op->else_body_)); - }(); - - // ── Helper: find the AssignStmt that defines a yield var in a branch ── - // The yield value is the SSA name of the diverging variable inside the - // branch. For tensor-typed phis, we need to resolve this to the - // tensor's actual source — the RHS of the assignment that sets it. - // This matters for the then-branch where ``boundary = zero`` creates a - // bare Python name that is NOT in the ``tensors`` dict, while - // ``zero`` (a kernel param) IS. - auto find_yield_source = [](const ir::StmtPtr& branch_body, const ir::VarPtr& yield_var) -> ir::ExprPtr { - const auto stmts = - ir::transform_utils::FlattenToStmts(ir::transform_utils::UnwrapAutoScope(branch_body)); - // Walk backwards (excluding the trailing YieldStmt) to find the - // most recent assignment to this var. - for (auto it = stmts.rbegin(); it != stmts.rend(); ++it) { - if (ir::As(*it)) continue; - if (auto assign = ir::As(*it)) { - if (assign->var_ == yield_var) return assign->value_; - } - } - return nullptr; - }; - - // ── Find a tensor init for tensor-typed phis ──────────────────── - // Tensor phis need an in-scope init to reference in the tensors dict. - // Scan function params first (guaranteed in tensors dict at runtime), - // then fall back to the then-branch yield's first tensor value. - std::string tensor_phi_init; - for (const auto& param : current_func_->params_) { - if (ir::AsTensorTypeLike(param->GetType())) { - tensor_phi_init = SanitizeName(param->name_hint_); - break; - } - } - if (tensor_phi_init.empty() && then_yield) { - for (const auto& val : then_yield->value_) { - const auto var = ir::As(val); - if (var && ir::AsTensorTypeLike(var->GetType())) { - auto src = find_yield_source(op->then_body_, var); - // Only resolve Var→Var aliases for init discovery. Call RHS - // triggers EmitCallToWorker which would submit a chip task - // before the `if` condition is evaluated. - if (src && ir::As(src) != nullptr) { - VisitExpr(src); - tensor_phi_init = current_expr_value_; - current_expr_value_ = ""; - } - break; - } - } - } + // ConvertToSSA always appends an else carrying the phi's incoming values + // (convert_to_ssa_pass.cpp:936-939), so a phi is always defined on both paths. + INTERNAL_CHECK_SPAN(op->return_vars_.empty() || op->else_body_.has_value(), op->span_) + << "Internal error: IfStmt with return_vars_ must carry an else_body " + "holding the phi's incoming values"; + + VisitExpr(op->condition_); + const std::string condition = current_expr_value_; + current_expr_value_ = ""; - // ── Pre-declare phi variables before the ``if`` ───────────────── - // Derive each tensor phi's initializer from its own yield source in - // the then-branch (preserving shape/dtype), falling back to the - // shared param-derived init or a zero-placeholder. - for (size_t i = 0; i < op->return_vars_.size(); ++i) { - const std::string phi_name = SanitizeName(op->return_vars_[i]->name_hint_); + // Merge each branch's yield into the phi name so post-if consumers see one + // name instead of a branch-local SSA name (issue #2180). + auto emit_yields = [&](const ir::StmtPtr& body) { + const auto yld = ir::transform_utils::GetLastYieldStmt(ir::transform_utils::UnwrapAutoScope(body)); + if (!yld) return; + for (size_t i = 0; i < op->return_vars_.size() && i < yld->value_.size(); ++i) { + VisitExpr(yld->value_[i]); + const std::string val = current_expr_value_; + current_expr_value_ = ""; + const std::string phi = SanitizeName(op->return_vars_[i]->name_hint_); if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) { - // Try a per-phi init from the then-branch yield first. - std::string per_phi_init; - if (then_yield && i < then_yield->value_.size()) { - const auto yield_var = ir::As(then_yield->value_[i]); - if (yield_var) { - auto src = find_yield_source(op->then_body_, yield_var); - // Only visit Var→Var aliases (e.g. ``boundary = zero``). - // For Call RHS, EmitCallToWorker would submit a chip task - // before the `if` condition — skip init discovery and fall - // through to the shared tensor_phi_init or zero placeholder. - if (src && ir::As(src) != nullptr) { - VisitExpr(src); - per_phi_init = current_expr_value_; - current_expr_value_ = ""; - } - } - } - if (!per_phi_init.empty()) { - emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + per_phi_init + "\"]"); - } else if (!tensor_phi_init.empty()) { - emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + tensor_phi_init + "\"]"); - } else { - emitter_.EmitLine("if \"" + phi_name + "\" not in tensors:"); - emitter_.EmitLine(" tensors[\"" + phi_name + - "\"] = torch.zeros((1,), dtype=torch.float32).share_memory_()"); - } + emitter_.EmitLine("tensors[\"" + phi + "\"] = tensors[\"" + val + "\"]"); } else { - emitter_.EmitLine(phi_name + " = None"); + emitter_.EmitLine(phi + " = " + val); } - declared_vars_.insert(phi_name); + declared_vars_.insert(phi); } - - // ── Emit below-branch yield-to-phi assignments ───────────────── - // Use the assignment RHS as the yield source when the RHS is a Var - // (e.g. ``boundary = zero`` — a kernel-param alias). For Call RHS - // (e.g. ``boundary = self.chip_run(...)``), EmitCallToWorker already - // places the result in ``tensors[yield_var_name]``, so use the yield - // var name directly. - auto emit_branch_yields = [&](const ir::StmtPtr& branch_body, const ir::YieldStmtPtr& yld) { - if (!yld) return; - for (size_t i = 0; i < op->return_vars_.size() && i < yld->value_.size(); ++i) { - const auto yield_var = ir::As(yld->value_[i]); - if (!yield_var) continue; - - // Find the assignment that defines this yield var. - auto src = find_yield_source(branch_body, yield_var); - // Only trace Var→Var aliases. For Call RHS the result is already - // in ``tensors[yield_var_name]`` via EmitCallToWorker. - const bool is_var_alias = src && ir::As(src) != nullptr; - const ir::ExprPtr yield_expr = is_var_alias ? src : yield_var; - VisitExpr(yield_expr); - const std::string yield_val = current_expr_value_; - current_expr_value_ = ""; - const std::string phi_name = SanitizeName(op->return_vars_[i]->name_hint_); - if (ir::AsTensorTypeLike(op->return_vars_[i]->GetType())) { - emitter_.EmitLine("tensors[\"" + phi_name + "\"] = tensors[\"" + yield_val + "\"]"); - } else { - emitter_.EmitLine(phi_name + " = " + yield_val); - } - } - }; - // ── Emit condition and if/else with yield-to-phi ──────────────── - VisitExpr(op->condition_); - std::string condition = current_expr_value_; - current_expr_value_ = ""; - - emitter_.EmitLine("if " + condition + ":"); - emitter_.IncreaseIndent(); - VisitStmt(op->then_body_); - emit_branch_yields(op->then_body_, then_yield); - emitter_.DecreaseIndent(); - - if (op->else_body_.has_value()) { - emitter_.EmitLine("else:"); - emitter_.IncreaseIndent(); - VisitStmt(*op->else_body_); - emit_branch_yields(*op->else_body_, else_yield); - emitter_.DecreaseIndent(); - } - return; - } - - // ── Simple if/else without return_vars_ (no phi nodes) ─────────── - VisitExpr(op->condition_); - std::string condition = current_expr_value_; - current_expr_value_ = ""; + }; emitter_.EmitLine("if " + condition + ":"); emitter_.IncreaseIndent(); VisitStmt(op->then_body_); + emit_yields(op->then_body_); emitter_.DecreaseIndent(); if (op->else_body_.has_value()) { emitter_.EmitLine("else:"); emitter_.IncreaseIndent(); VisitStmt(*op->else_body_); + emit_yields(*op->else_body_); emitter_.DecreaseIndent(); } } diff --git a/tests/ut/codegen/distributed/test_host_orch_distributed.py b/tests/ut/codegen/distributed/test_host_orch_distributed.py index 5fb1777271..8a868c1d93 100644 --- a/tests/ut/codegen/distributed/test_host_orch_distributed.py +++ b/tests/ut/codegen/distributed/test_host_orch_distributed.py @@ -956,9 +956,9 @@ def test_host_collective_builtin_template_package_exists(package_name, variant): # --------------------------------------------------------------------------- -def test_if_cross_branch_phi_predeclares_and_yields_tensors() -> None: +def test_if_cross_branch_phi_yields_tensors() -> None: """Tensor phi emitted by ConvertToSSA for a cross-branch diverging variable - is pre-declared and yielded via ``tensors[...]``, not bare Python names.""" + is yielded via ``tensors[...]`` assignments, not bare Python names (issue #2180).""" SIZE = 4 P = 2 @@ -986,41 +986,45 @@ def host_orch( code = _lower(Prog, convert_to_ssa=True) - # Pre-declared phi before the ``if`` — must use tensors[...] for tensor phis. - pre_if = code.split("if", 1)[0] - assert re.search(r'tensors\["boundary__phi_v\d+"\]\s*=.*tensors', pre_if), ( - "Tensor phi must be pre-declared via tensors[...] before if block:\n" + code + # Sanity check: the generated Python must be syntactically valid. + compile(code, "", "exec") + + # Locate the ``if`` line via line-anchored regex (avoids false matches + # on substrings within comments, identifiers, or guard lines). + lines = code.splitlines() + if_idx = None + for i, line in enumerate(lines): + if re.search(r"^\s*if\s+.*:\s*$", line): + if_idx = i + break + assert if_idx is not None, f"No if statement found in generated code:\n{code}" + + # Find the matching ``else:`` at the same indent level. + if_indent = len(lines[if_idx]) - len(lines[if_idx].lstrip()) + else_idx = None + for i in range(if_idx + 1, len(lines)): + stripped = lines[i].lstrip() + if stripped.startswith("else:") and (len(lines[i]) - len(stripped)) == if_indent: + else_idx = i + break + assert else_idx is not None, f"No else at indent level {if_indent} in generated code:\n{code}" + + then_block = "\n".join(lines[if_idx + 1 : else_idx]) + else_block = "\n".join(lines[else_idx + 1 :]) + + # Both branches must emit a tensors-based phi assignment so the merged + # scope has a single ``tensors["boundary__phi_v<...>"]`` entry. + phi_yield_pat = r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]' + assert re.search(phi_yield_pat, then_block), ( + "Then-branch must yield tensor phi via tensors[...] = tensors[...]:\n" + code ) - - # Isolate the then-branch body: everything after the ``if ...:`` line header - # (first newline after the ``if`` keyword) up to the ``else:`` boundary. - if "else:" in code: - post_if = code.split("if", 1)[1].split("\n", 1)[1] # strip ``:\n`` header - then_block = post_if.split("else:")[0] - else_block = post_if.split("else:")[1] - else: - post_if = code.split("if", 1)[1].split("\n", 1)[1] - then_block = post_if - else_block = "" - - # Yield in the then-branch must reference a valid tensors-dict entry (no - # bare Python names for tensor phis). - then_yield = re.search( - r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', - then_block, - ) - assert then_yield is not None, "Then-branch must yield phi via tensors[...] = tensors[...]:\n" + code - - # Yield in the else-branch likewise. - else_yield = re.search( - r'tensors\["boundary__phi_v\d+"\]\s*=\s*tensors\["[^"]+"\]', - else_block, + assert re.search(phi_yield_pat, else_block), ( + "Else-branch must yield tensor phi via tensors[...] = tensors[...]:\n" + code ) - assert else_yield is not None, "Else-branch must yield phi via tensors[...] = tensors[...]:\n" + code # No bare-Python-name tensor assignment in the then-branch (the bug was # ``boundary__ssa_v0 = zero__ssa_v0`` which raised NameError). - assert not re.search(r"(?