Introduce ActBeforeDifferentiatingLoopCondition to fix the error esti…

…mation mode.
vgvassilev · Jan 6, 2024 · dea82f4 · dea82f4
1 parent 6904f1a
commit dea82f4
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 4 deletions.
diff --git a/include/clad/Differentiator/ErrorEstimator.h b/include/clad/Differentiator/ErrorEstimator.h
@@ -255,6 +255,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   void ActOnEndOfDerivedFnBody() override;
   void ActBeforeDifferentiatingStmtInVisitCompoundStmt() override;
   void ActAfterProcessingStmtInVisitCompoundStmt() override;
+  void ActBeforeDifferentiatingLoopCondition() override;
   void ActBeforeDifferentiatingSingleStmtBranchInVisitIfStmt() override;
   void ActBeforeFinalisingVisitBranchSingleStmtInIfVisitStmt() override;
   void ActBeforeDifferentiatingLoopInitStmt() override;

diff --git a/include/clad/Differentiator/ExternalRMVSource.h b/include/clad/Differentiator/ExternalRMVSource.h
@@ -106,6 +106,9 @@ class ExternalRMVSource {
   /// branch in `VisitBranch` lambda in
   virtual void ActBeforeFinalisingVisitBranchSingleStmtInIfVisitStmt() {}
 
+  /// This is called just before differentiating loop conditions.
+  virtual void ActBeforeDifferentiatingLoopCondition() {}
+
   /// This is called just before differentiating init statement of loops.
   virtual void ActBeforeDifferentiatingLoopInitStmt() {}
 

diff --git a/include/clad/Differentiator/MultiplexExternalRMVSource.h b/include/clad/Differentiator/MultiplexExternalRMVSource.h
@@ -41,6 +41,7 @@ class MultiplexExternalRMVSource : public ExternalRMVSource {
   void ActAfterProcessingStmtInVisitCompoundStmt() override;
   void ActBeforeDifferentiatingSingleStmtBranchInVisitIfStmt() override;
   void ActBeforeFinalisingVisitBranchSingleStmtInIfVisitStmt() override;
+  void ActBeforeDifferentiatingLoopCondition() override;
   void ActBeforeDifferentiatingLoopInitStmt() override;
   void ActBeforeDifferentiatingSingleStmtLoopBody() override;
   void ActAfterProcessingSingleStmtBodyInVisitForLoop() override;

diff --git a/lib/Differentiator/ErrorEstimator.cpp b/lib/Differentiator/ErrorEstimator.cpp
@@ -562,6 +562,11 @@ void ErrorEstimationHandler::ActAfterProcessingStmtInVisitCompoundStmt() {
   EmitErrorEstimationStmts(direction::reverse);
 }
 
+void ErrorEstimationHandler::
+    ActBeforeDifferentiatingLoopCondition() {
+  m_ShouldEmit.push(true);
+}
+
 void ErrorEstimationHandler::
     ActBeforeDifferentiatingSingleStmtBranchInVisitIfStmt() {
   m_ShouldEmit.push(true);

diff --git a/lib/Differentiator/MultiplexExternalRMVSource.cpp b/lib/Differentiator/MultiplexExternalRMVSource.cpp
@@ -120,6 +120,13 @@ void MultiplexExternalRMVSource::
   }
 }
 
+void MultiplexExternalRMVSource::
+    ActBeforeDifferentiatingLoopCondition() {
+  for (auto source : m_Sources) {
+    source->ActBeforeDifferentiatingLoopCondition();
+  }
+}
+
 void MultiplexExternalRMVSource::ActBeforeDifferentiatingLoopInitStmt() {
   for (auto source : m_Sources) {
     source->ActBeforeDifferentiatingLoopInitStmt();

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -3273,8 +3273,11 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     StmtDiff bodyDiff = nullptr;
     StmtDiff condDiff = nullptr;
     if (isa<CompoundStmt>(body)) {
-      if (cond)
+      if (cond) {
+        if (m_ExternalSource)
+          m_ExternalSource->ActBeforeDifferentiatingLoopCondition();
         condDiff = DifferentiateSingleStmt(cond, /*dfdS=*/nullptr);
+      }
       bodyDiff = Visit(body);
 
       bodyDiff.updateStmt(utils::PrependAndCreateCompoundStmt(
@@ -3287,8 +3290,12 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
       beginScope(Scope::DeclScope);
       beginBlock(direction::forward);
       addToCurrentBlock(counterIncrement);
-      if (cond)
+
+      if (cond) {
+        if (m_ExternalSource)
+          m_ExternalSource->ActBeforeDifferentiatingLoopCondition();
         condDiff = DifferentiateSingleStmt(cond, /*dfdS=*/nullptr);
+      }
       if (m_ExternalSource)
         m_ExternalSource->ActBeforeDifferentiatingSingleStmtLoopBody();
       bodyDiff = DifferentiateSingleStmt(body, /*dfdS=*/nullptr);
@@ -3338,10 +3345,9 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
 
     // `for` loops have counter decrement expression in the
     // loop iteration-expression.
-    if (!isForLoop) {
+    if (!isForLoop)
       bodyDiff.updateStmtDx(utils::AppendAndCreateCompoundStmt(
               m_Context, bodyDiff.getStmt_dx(), counterDecrement));
-    }
     return {bodyDiff, condDiff};
   }