Add shouldUseAtomicOps function

vgvassilev · Oct 3, 2024 · 873ded3 · 873ded3
1 parent 5c0a26d
commit 873ded3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h
@@ -434,6 +434,10 @@ namespace clad {
       return StmtDiff{Clone(NS), Clone(NS)};
     };
 
+    /// Helper function that checks whether the function to be derived
+    /// is meant to be executed only by the GPU
+    bool shouldUseCudaAtomicOps();
+
     /// Add call to cuda::atomicAdd for the given LHS and RHS expressions.
     ///
     /// \param[in] LHS The left-hand side expression.

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -104,6 +104,12 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     return CladTapeResult{*this, PushExpr, PopExpr, TapeRef};
   }
 
+  bool ReverseModeVisitor::shouldUseCudaAtomicOps() {
+    return m_DiffReq->hasAttr<clang::CUDAGlobalAttr>() ||
+           (m_DiffReq->hasAttr<clang::CUDADeviceAttr>() &&
+            !m_DiffReq->hasAttr<clang::CUDAHostAttr>());
+  }
+
   clang::Expr* ReverseModeVisitor::BuildCallToCudaAtomicAdd(clang::Expr* LHS,
                                                             clang::Expr* RHS) {
     DeclarationName atomicAddId = &m_Context.Idents.get("atomicAdd");
@@ -1518,9 +1524,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
         BuildArraySubscript(target, forwSweepDerivativeIndices);
     // Create the (target += dfdx) statement.
     if (dfdx()) {
-      if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>() ||
-          (m_DiffReq->hasAttr<clang::CUDADeviceAttr>() &&
-           !m_DiffReq->hasAttr<clang::CUDAHostAttr>())) {
+      if (shouldUseCudaAtomicOps()) {
         Expr* atomicCall = BuildCallToCudaAtomicAdd(result, dfdx());
         // Add it to the body statements.
         addToCurrentBlock(atomicCall, direction::reverse);
@@ -2320,9 +2324,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
         derivedE = BuildOp(UnaryOperatorKind::UO_Deref, diff_dx);
         // Create the (target += dfdx) statement.
         if (dfdx()) {
-          if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>() ||
-              (m_DiffReq->hasAttr<clang::CUDADeviceAttr>() &&
-               !m_DiffReq->hasAttr<clang::CUDAHostAttr>())) {
+          if (shouldUseCudaAtomicOps()) {
             Expr* atomicCall = BuildCallToCudaAtomicAdd(diff_dx, dfdx());
             // Add it to the body statements.
             addToCurrentBlock(atomicCall, direction::reverse);