Recompute the values of CUDA built-in index functions

Author: ovdiiuv

include/clad/Differentiator/ReverseModeVisitor.h

@@ -210,6 +210,8 @@ 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.
+    bool isCUDABuiltInIndex(const clang::Expr* E);
 
     /// Builds a variable declaration and stores it in the function
     /// global scope.

lib/Differentiator/ReverseModeVisitor.cpp

@@ -3228,9 +3228,31 @@ 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();
+                        if (const auto* DRE = llvm::dyn_cast<DeclRefExpr>(innerE)) {
+                            std::string Name = DRE->getNameInfo().getAsString();
+                            return Name == "blockDim" || Name == "threadIdx" || Name == "blockIdx";
+                        }
+                    }
+                }
+            }
+        }
+    }
+    return false;
+  }
+  
   bool ReverseModeVisitor::UsefulToStoreGlobal(Expr* E) {
     if (!E)
       return false;