diff --git a/src/codegen/distributed/distributed_codegen.cpp b/src/codegen/distributed/distributed_codegen.cpp index f418204129..965875457f 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,23 @@ 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). + // 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_ + "\"]"); + declared_vars_.insert(var_name); + current_expr_value_ = ""; + } + current_target_var_ = ""; + return; + } + // Standard expression VisitExpr(op->value_); @@ -813,19 +831,46 @@ void DistributedCodegen::VisitStmt_(const ir::ForStmtPtr& op) { void DistributedCodegen::VisitStmt_(const ir::IfStmtPtr& op) { INTERNAL_CHECK(op != nullptr) << "Internal error: null IfStmt"; + // 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_); - std::string condition = current_expr_value_; + const std::string condition = current_expr_value_; current_expr_value_ = ""; + // 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())) { + emitter_.EmitLine("tensors[\"" + phi + "\"] = tensors[\"" + val + "\"]"); + } else { + emitter_.EmitLine(phi + " = " + val); + } + declared_vars_.insert(phi); + } + }; + 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 9219425283..8a868c1d93 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,83 @@ 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_yields_tensors() -> None: + """Tensor phi emitted by ConvertToSSA for a cross-branch diverging variable + is yielded via ``tensors[...]`` assignments, not bare Python names (issue #2180).""" + 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 # pyright: ignore[reportReturnType] + + code = _lower(Prog, convert_to_ssa=True) + + # 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 + ) + assert re.search(phi_yield_pat, else_block), ( + "Else-branch must yield tensor 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"boundary__ssa_v\d+\s*=.*\w+__ssa", then_block), ( + "Then-branch must not contain bare Python tensor assignments:\n" + code + ) + + if __name__ == "__main__": pytest.main([__file__, "-v"])