Skip to content

Lower conditional merges to predecessor copies #1654

Description

@c0d1f1ed

Background

Warp lowers variables merged across if/else, including conditional expressions, to a post-branch select:

var_merged = wp::where(cond, if_var, else_var);

where reads only the selected arm, whose defining branch matches cond, so ordinary merges are well-defined. However, the generated target code leaves each operand initialized only along its correlated branch, a backend pattern worth avoiding.

GH-1574 exposed a more serious nested form: a symbol-map leak referenced an if-branch value inside its sibling else, allowing an intermediate merge to select an uninitialized value. Under register pressure this surfaced as an illegal CUDA memory access. MR !2541 fixed that bug; this follow-up separately hardens ordinary merges.

Proposal

Replace the merge-point select with standard SSA destruction copies in predecessor blocks:

wp::float32 var_merged;
if (var_5)  { …; var_merged = wp::copy(var_10); }
if (!var_5) { …; var_merged = wp::copy(var_15); }

Validation

  • Preserve reverse-mode gradients despite assigning one destination in multiple branches, which relaxes the one-definition-per-Var assumption.
  • Verify forward/replay copies and guarded adjoint propagation in generated source, with no post-branch merge over conditionally defined operands.
  • Spot-check representative generated code for a performance regression.

Scope

Update Adjoint.emit_If, Adjoint.emit_IfExp, and their reverse-mode
lowering. Treat this as a dedicated codegen-hardening change.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Fields

No fields configured for Enhancement.

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions