Recompute the values of CUDA built-in index functions

Author: ovdiiuv

include/clad/Differentiator/ReverseModeVisitor.h

@@ -210,6 +210,10 @@ namespace clad {
     /// For an expr E, decides if we should recompute it or store it.
     /// This is the central point for checkpointing.
     bool ShouldRecompute(const clang::Expr* E);
+    /// Called in ShouldRecompute. In CUDA, to access a current thread/block id
+    /// we use functions that do not change the state of any variable, since no
+    /// point to store the value.
+    static bool isCUDABuiltInIndex(const clang::Expr* E);
 
     /// Builds a variable declaration and stores it in the function
     /// global scope.

lib/Differentiator/ReverseModeVisitor.cpp

@@ -3228,7 +3228,30 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
   }
 
   bool ReverseModeVisitor::ShouldRecompute(const Expr* E) {
-    return !(utils::ContainsFunctionCalls(E) || E->HasSideEffects(m_Context));
+    return !(utils::ContainsFunctionCalls(E) || E->HasSideEffects(m_Context)) ||
+           isCUDABuiltInIndex(E);
+  }
+
+  bool ReverseModeVisitor::isCUDABuiltInIndex(const Expr* E) {
+    const clang::Expr* B = E->IgnoreImplicit();
+    if (const auto* pseudoE = llvm::dyn_cast<PseudoObjectExpr>(B)) {
+      for (const Expr* subE : pseudoE->semantics()) {
+        subE = subE->IgnoreImplicit();
+        if (const auto* call = llvm::dyn_cast<CallExpr>(subE)) {
+          if (const auto* memberE = llvm::dyn_cast<MemberExpr>(
+                  call->getCallee()->IgnoreParenImpCasts())) {
+            const Expr* base = memberE->getBase()->IgnoreImplicit();
+            if (const auto* opaqueE = llvm::dyn_cast<OpaqueValueExpr>(base)) {
+              const Expr* innerE = opaqueE->getSourceExpr()->IgnoreImplicit();
+              QualType innerT = innerE->getType();
+              if (innerT.isConstQualified())
+                return true;
+            }
+          }
+        }
+      }
+    }
+    return false;
   }
 
   bool ReverseModeVisitor::UsefulToStoreGlobal(Expr* E) {