Ensure proper ending of the reverse block for conditional statements. F…

…ixes vgvassilev#922.
PetroZarytskyi · Jun 7, 2024 · 7d330c3 · 7d330c3
1 parent 3cf8b33
commit 7d330c3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -945,20 +945,20 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
                          .get();
     // If result is a glvalue, we should keep it as it can potentially be
     // assigned as in (c ? a : b) = x;
+    Expr* ResultRef = nullptr;
     if ((CO->isModifiableLvalue(m_Context) == Expr::MLV_Valid) &&
         ifTrueExprDiff.getExpr_dx() && ifFalseExprDiff.getExpr_dx()) {
-      Expr* ResultRef = m_Sema
-                            .ActOnConditionalOp(noLoc, noLoc, condStored,
-                                                ifTrueExprDiff.getExpr_dx(),
-                                                ifFalseExprDiff.getExpr_dx())
-                            .get();
+      ResultRef = m_Sema
+                      .ActOnConditionalOp(noLoc, noLoc, condStored,
+                                          ifTrueExprDiff.getExpr_dx(),
+                                          ifFalseExprDiff.getExpr_dx())
+                      .get();
       if (ResultRef->isModifiableLvalue(m_Context) != Expr::MLV_Valid)
         ResultRef = nullptr;
-      Stmt* revBlock = utils::unwrapIfSingleStmt(endBlock(direction::reverse));
-      addToCurrentBlock(revBlock, direction::reverse);
-      return StmtDiff(condExpr, ResultRef);
     }
-    return StmtDiff(condExpr);
+    Stmt* revBlock = utils::unwrapIfSingleStmt(endBlock(direction::reverse));
+    addToCurrentBlock(revBlock, direction::reverse);
+    return StmtDiff(condExpr, ResultRef);
   }
 
   StmtDiff ReverseModeVisitor::VisitForStmt(const ForStmt* FS) {

diff --git a/test/Gradient/Gradients.C b/test/Gradient/Gradients.C
@@ -913,6 +913,18 @@ double fn_null_stmts(double x) {
 //CHECK-NEXT:              *_d_x += 1;
 //CHECK-NEXT:          }
 
+double fn_const_cond_op(double x) {
+  return x + (x > 0 ? 1.0 : 0.0);
+}
+
+//CHECK:               void fn_const_cond_op_grad(double x, double *_d_x) {
+//CHECK-NEXT:              bool _cond0;
+//CHECK-NEXT:              _cond0 = x > 0;
+//CHECK-NEXT:              goto _label0;
+//CHECK-NEXT:            _label0:
+//CHECK-NEXT:              *_d_x += 1;
+//CHECK-NEXT:          }
+
 #define TEST(F, x, y)                                                          \
   {                                                                            \
     result[0] = 0;                                                             \
@@ -985,4 +997,7 @@ int main() {
   TEST_GRADIENT(fn_cond_init, /*numOfDerivativeArgs=*/1, -1, &dx); // CHECK-EXEC: 1.00
 
   INIT_GRADIENT(fn_null_stmts);
+
+  INIT_GRADIENT(fn_const_cond_op);
+  TEST_GRADIENT(fn_const_cond_op, /*numOfDerivativeArgs=*/1, 0, &dx); // CHECK-EXEC: 1.00
 }