Add reverse mode support for omp parallel directive

Author: Shubham Shukla

include/clad/Differentiator/ReverseModeVisitor.h

@@ -481,6 +481,7 @@ namespace clad {
     StmtDiff
     VisitOMPParallelForDirective(const clang::OMPParallelForDirective* D);
     StmtDiff VisitOMPCriticalDirective(const clang::OMPCriticalDirective* D);
+    StmtDiff VisitOMPParallelDirective(const clang::OMPParallelDirective* D);
 
     /// Helper function that builds `T* _this = malloc(sifeof(T));`
     /// and `free(_this)`.

lib/Differentiator/ReverseModeVisitorOpenMP.cpp

@@ -521,8 +521,8 @@ StmtDiff ReverseModeVisitor::VisitOMPExecutableDirective(
     llvm::SaveAndRestore<bool> SaveisInsideOMPBlock(isInsideOMPBlock);
     isInsideOMPBlock = true;
 
-    CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(OMPD_parallel,
-                                                                  nullptr);
+    CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(
+        OMPD_parallel, getCurrentScope());
     StmtDiff BodyDiff;
     {
       Sema::CompoundScopeRAII CompoundScope(m_Sema);
@@ -537,8 +537,8 @@ StmtDiff ReverseModeVisitor::VisitOMPExecutableDirective(
                         .ActOnOpenMPRegionEnd(BodyDiff.getStmt(), OrigClauses)
                         .get();
 
-    CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(OMPD_parallel,
-                                                                  nullptr);
+    CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(
+        OMPD_parallel, getCurrentScope());
     // Visit twice, but use only the result of the first visit, for capture
     // variables only.
     {
@@ -549,7 +549,17 @@ StmtDiff ReverseModeVisitor::VisitOMPExecutableDirective(
         const auto* FS = cast<ForStmt>(CS);
         DifferentiateCanonicalLoop(FS);
       } else {
-        Visit(CS);
+        StmtDiff SecondDiff = Visit(CS);
+        auto& II = m_Context.Idents.get("_clad_reverse_guard");
+        DeclarationNameInfo CritDirName(DeclarationName(&II), D->getBeginLoc());
+        llvm::SmallVector<OMPClause*, 0> NoClauses;
+        Stmt* ProtectedReverse =
+            CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema)
+                .ActOnOpenMPExecutableDirective(
+                    OMPD_critical, CritDirName, OMPD_unknown, NoClauses,
+                    SecondDiff.getStmt_dx(), D->getBeginLoc(), D->getEndLoc())
+                .get();
+        BodyDiff = {BodyDiff.getStmt(), ProtectedReverse};
       }
       m_Globals.swap(temp);
     }
@@ -574,6 +584,15 @@ StmtDiff ReverseModeVisitor::VisitOMPExecutableDirective(
                                               D->getBeginLoc(), D->getEndLoc())
               .get()};
 }
+StmtDiff
+ReverseModeVisitor::VisitOMPParallelDirective(const OMPParallelDirective* D) {
+  DeclarationNameInfo DirName;
+  CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).StartOpenMPDSABlock(
+      OMPD_parallel, DirName, nullptr, D->getBeginLoc());
+  StmtDiff SDiff = VisitOMPExecutableDirective(D);
+  CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).EndOpenMPDSABlock(SDiff.getStmt());
+  return SDiff;
+}
 StmtDiff ReverseModeVisitor::VisitOMPParallelForDirective(
     const OMPParallelForDirective* D) {
   DeclarationNameInfo DirName;

test/Gradient/OpenMP.C

@@ -944,6 +944,48 @@ double fn23(const double *x, int n) {
 // CHECK-NEXT:         }
 // CHECK-NEXT: }
 
+double fn24(const double *x, int n) {
+    double result = 0;
+    #pragma omp parallel reduction(+:result)
+    {
+        result += x[0] * x[1];
+        result += x[1] * x[2];
+        result += x[2] * x[3];
+    }
+    return result;
+}
+
+// CHECK:  void fn24_grad(const double *x, int n, double *_d_x, int *_d_n) {
+// CHECK-NEXT:      double _d_result = 0.;
+// CHECK-NEXT:      double result = 0;
+// CHECK-NEXT:      #pragma omp parallel reduction(+: result)
+// CHECK-NEXT:          {
+// CHECK-NEXT:              result += x[0] * x[1];
+// CHECK-NEXT:              result += x[1] * x[2];
+// CHECK-NEXT:              result += x[2] * x[3];
+// CHECK-NEXT:          }
+// CHECK-NEXT:      _d_result += 1;
+// CHECK-NEXT:      #pragma omp parallel private(result) firstprivate(_d_result)
+// CHECK-NEXT:          #pragma omp critical (_clad_reverse_guard)
+// CHECK-NEXT:              {
+// CHECK-NEXT:                  {
+// CHECK-NEXT:                      double _r_d5 = _d_result;
+// CHECK-NEXT:                      _d_x[2] += _r_d5 * x[3];
+// CHECK-NEXT:                      _d_x[3] += x[2] * _r_d5;
+// CHECK-NEXT:                  }
+// CHECK-NEXT:                  {
+// CHECK-NEXT:                      double _r_d4 = _d_result;
+// CHECK-NEXT:                      _d_x[1] += _r_d4 * x[2];
+// CHECK-NEXT:                      _d_x[2] += x[1] * _r_d4;
+// CHECK-NEXT:                  }
+// CHECK-NEXT:                  {
+// CHECK-NEXT:                      double _r_d3 = _d_result;
+// CHECK-NEXT:                      _d_x[0] += _r_d3 * x[1];
+// CHECK-NEXT:                      _d_x[1] += x[0] * _r_d3;
+// CHECK-NEXT:                  }
+// CHECK-NEXT:              }
+// CHECK-NEXT:  }
+
 template <size_t N>
 void reset(double (&arr)[N], double val = 0) {
   for (size_t i = 0; i < N; ++i)
@@ -1067,5 +1109,11 @@ int main() {
   auto fn23_grad = clad::gradient(fn23);
   fn23_grad.execute(x, 4, dx, &dn);
   printf("{%.2f, %.2f, %.2f, %.2f}\n", dx[0], dx[1], dx[2], dx[3]); // CHECK-EXEC: {4.00, 6.00, 8.00, 10.00}
+
+  omp_set_num_threads(2);
+  reset(dx);
+  auto fn24_grad = clad::gradient(fn24);
+  fn24_grad.execute(x, 4, dx, &dn);
+  printf("{%.2f, %.2f, %.2f, %.2f}\n", dx[0], dx[1], dx[2], dx[3]); // CHECK-EXEC: {6.00, 12.00, 16.00, 8.00}
   return 0;
 }