fix(codegen): emit tensor phis via tensors dict to avoid NameError (#2180)

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<ir::TensorType>(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();
   }
 }

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, "<host_orch>", "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"])