Recompute the values of CUDA built-in index functions

Author: ovdiiuv

lib/Differentiator/ReverseModeVisitor.cpp

@@ -3232,8 +3232,26 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     return Visit(EWC->getSubExpr(), dfdx());
   }
 
+  /// 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 Expr* E) {
+    const clang::Expr* B = E->IgnoreImplicit();
+    if (const auto* pseudoE = llvm::dyn_cast<PseudoObjectExpr>(B)) {
+      if (const auto* opaqueE =
+              llvm::dyn_cast<OpaqueValueExpr>(pseudoE->getSemanticExpr(0))) {
+        const Expr* innerE = opaqueE->getSourceExpr()->IgnoreImplicit();
+        QualType innerT = innerE->getType();
+        if (innerT.isConstQualified())
+          return true;
+      }
+    }
+    return false;
+  }
+
   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::UsefulToStoreGlobal(Expr* E) {