Add reverse-mode support for CUDA shared memory attributes

include/clad/Differentiator/BuiltinDerivativesCUDA.cuh

@@ -2,6 +2,11 @@
 
 namespace clad {
 
+__device__ inline unsigned int get_dynamic_smem_size() {
+  unsigned int smem_size;
+  asm volatile("mov.u32 %0, %%dynamic_smem_size;" : "=r"(smem_size));
+  return smem_size;
+}
 namespace custom_derivatives {
 
 __device__ inline void __expf_pullback(float a, float d_y, float* d_a) {

include/clad/Differentiator/ReverseModeVisitor.h

@@ -802,6 +802,10 @@ namespace clad {
 
     /// A flag indicating if the Stmt is contained in a checkpointed loop.
     bool m_IsInsideCheckpointedLoop = false;
+    void HandleCUDASharedMemoryDecl(
+        const clang::VarDecl* VD, clang::VarDecl* VDForward,
+        clang::VarDecl* VDDerived,
+        llvm::SmallVectorImpl<clang::Stmt*>& memsetCalls);
   };
 } // end namespace clad
 

include/clad/Differentiator/VisitorBase.h

@@ -378,7 +378,8 @@ namespace clad {
                                        clang::Expr* Init = nullptr,
                                        bool DirectInit = false,
                                        clang::TypeSourceInfo* TSI = nullptr,
-                                       clang::StorageClass SC = clang::SC_None);
+                                       clang::StorageClass SC = clang::SC_None,
+                                       const clang::VarDecl* OrigVD = nullptr);
     /// Creates a namespace declaration and enters its context. All subsequent
     /// Stmts are built inside that namespace, until
     /// m_Sema.PopDeclContextIsUsed.

lib/Differentiator/CMakeLists.txt

@@ -65,6 +65,7 @@ llvm_add_library(cladDifferentiator
   PushForwardModeVisitor.cpp
   ReverseModeForwPassVisitor.cpp
   ReverseModeVisitor.cpp
+  ReverseModeVisitorCUDA.cpp
   ReverseModeVisitorOpenMP.cpp
   TBRAnalyzer.cpp
   Timers.cpp

lib/Differentiator/ReverseModeVisitor.cpp

@@ -151,6 +151,10 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     if (!m_Context.getLangOpts().CUDA)
       return false;
     if (const auto* DRE = dyn_cast<DeclRefExpr>(E)) {
+      if (const auto* VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+        if (VD->hasAttr<clang::CUDASharedAttr>())
+          return true;
+      }
       if (const auto* PVD = dyn_cast<ParmVarDecl>(DRE->getDecl())) {
         if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>())
           // Check whether this param is in the global memory of the GPU
@@ -168,6 +172,12 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     } else if (const auto* ASE = dyn_cast<ArraySubscriptExpr>(E)) {
       const auto* base =
           dyn_cast<DeclRefExpr>(ASE->getBase()->IgnoreImpCasts());
+      if (const auto* VD = dyn_cast<VarDecl>(base->getDecl())) {
+        if (VD->hasAttr<clang::CUDASharedAttr>()) {
+          const auto* idx = ASE->getIdx();
+          return !clad::utils::isInjective(idx, m_DiffReq.m_AnalysisDC);
+        }
+      }
       if (const auto* PVD = dyn_cast<ParmVarDecl>(base->getDecl())) {
         const auto* idx = ASE->getIdx();
         if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>())
@@ -314,8 +324,14 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     DeclarationNameInfo DNI = utils::BuildDeclarationNameInfo(m_Sema, name);
     DeclWithContext result = m_Builder.cloneFunction(m_DiffReq.Function, *this,
                                                      DC, loc, DNI, dFnType);
+
     m_Derivative = result.first;
 
+    if (m_DiffReq.Function->hasAttr<clang::CUDAGlobalAttr>()) {
+      auto* GlobalAtt = clang::CUDAGlobalAttr::CreateImplicit(m_Context);
+      GlobalAtt->setImplicit(false);
+      m_Derivative->addAttr(GlobalAtt);
+    }
     // Function declaration scope
     beginScope(Scope::FunctionPrototypeScope | Scope::FunctionDeclarationScope |
                Scope::DeclScope);
@@ -1958,6 +1974,10 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     // simplest way to support begin/end functions of the former and not deal
     // with the type mismatch.
     std::string FDName = FD->getNameAsString();
+    if (FDName == "__syncthreads") {
+      addToCurrentBlock(Clone(CE), direction::reverse);
+      return StmtDiff(Clone(CE), nullptr);
+    }
     if (FDName == "begin" || FDName == "end") {
       const Expr* arg = nullptr;
       if (const auto* MCE = dyn_cast<CXXMemberCallExpr>(CE))
@@ -3052,18 +3072,41 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
         dummyInit = m_Sema.ActOnInitList(noLoc, args, noLoc).get();
       }
     }
+    bool isDynamicSharedMem =
+        VD->hasAttr<CUDASharedAttr>() && VD->getType()->isIncompleteArrayType();
+
+    if (isDynamicSharedMem) {
+      QualType ElemTy = cast<IncompleteArrayType>(VD->getType().getTypePtr())
+                            ->getElementType();
+      VDDerivedType = m_Context.getPointerType(ElemTy);
+      dummyInit = nullptr; // Now perfectly in scope!
+    }
 
     StorageClass SC = isInsideOMPBlock ? SC_Static : SC_None;
 
+    if (VD->getStorageClass() == clang::SC_Extern)
+      SC = clang::SC_Extern;
     // Build the adjoint VarDecl
     VarDecl* VDDerived = nullptr;
     if (m_DiffReq.shouldHaveAdjoint(VD) &&
         !clad::utils::hasNonDifferentiableAttribute(VD)) {
       if (!isLambdaDS) {
         llvm::StringRef Name = VD->getName();
         std::string CleanName = Name.ltrim('_').str();
-        VDDerived = BuildGlobalVarDecl(VDDerivedType, "_d_" + CleanName,
-                                       dummyInit, false, nullptr, SC);
+        if (isDynamicSharedMem) {
+          QualType ElemTy =
+              cast<IncompleteArrayType>(VD->getType().getTypePtr())
+                  ->getElementType();
+          QualType PointerTy = m_Context.getPointerType(ElemTy);
+          VDDerived = BuildGlobalVarDecl(PointerTy, "_d_" + CleanName,
+                                         /* dummyInit = */ nullptr, false,
+                                         nullptr, SC_None);
+        } else {
+          VDDerived = BuildGlobalVarDecl(VDDerivedType, "_d_" + CleanName,
+                                         dummyInit, false, nullptr, SC);
+        }
+        if (!isDynamicSharedMem && VD->hasAttr<clang::CUDASharedAttr>())
+          VDDerived->setInit(nullptr);
       }
     }
 
@@ -3160,19 +3203,53 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
                                  initDiff.getExpr(), VD->isDirectInit(),
                                  VDCloneTSI, SC);
 
+    if (isDynamicSharedMem && VDDerived) {
+
+      llvm::SmallVector<Expr*, 0> args;
+      Expr* sizeCall = GetFunctionCall("get_dynamic_smem_size", "clad", args);
+
+      Expr* two = ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context,
+                                                    /*val=*/2);
+      Expr* halfSize = BuildOp(BO_Div, sizeCall, two);
+
+      Expr* primalRef = BuildDeclRef(VDClone);
+      QualType charPtrTy = m_Context.getPointerType(m_Context.CharTy);
+      Expr* castPrimal =
+          m_Sema
+              .BuildCStyleCastExpr(
+                  noLoc, m_Context.getTrivialTypeSourceInfo(charPtrTy), noLoc,
+                  primalRef)
+              .get();
+
+      Expr* byteOffsetPtr = BuildOp(BO_Add, castPrimal, halfSize);
+      Expr* parenOffsetPtr = BuildParens(byteOffsetPtr);
+      Expr* finalInit =
+          m_Sema
+              .BuildCStyleCastExpr(
+                  noLoc, m_Context.getTrivialTypeSourceInfo(VDDerivedType),
+                  noLoc, parenOffsetPtr)
+              .get();
+
+      initDiff.updateStmtDx(finalInit);
+    }
+
     // The choice of isDirectInit is mostly stylistic.
     bool isRealConstArray = false;
     if (const auto* arrType = dyn_cast<ConstantArrayType>(VDType))
       isRealConstArray = arrType->getElementType()->isRealType();
     bool isDirectInit = VD->isDirectInit() && (!RD || isNonAggrClass);
-    if (VDDerivedType->isBuiltinType() || !VD->getInit() || isRealConstArray) {
+    if (!isDynamicSharedMem && (VDDerivedType->isBuiltinType() ||
+                                !VD->getInit() || isRealConstArray)) {
       initDiff.updateStmtDx(getZeroInit(VDType));
       isDirectInit = false;
     } else if (Expr* size = getStdInitListSizeExpr(VD->getInit())) {
       initDiff.updateStmtDx(Clone(size));
       isConstructInit = true;
     }
 
+    if (!isDynamicSharedMem && VD->hasAttr<clang::CUDASharedAttr>())
+      initDiff.updateStmtDx(nullptr);
+
     // Update the initializer
     if (VDDerived)
       SetDeclInit(VDDerived, initDiff.getExpr_dx(), isDirectInit);
@@ -3337,10 +3414,17 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
           else {
             VarDecl* VDDerived = VDDiff.getDecl_dx();
             declsDiff.push_back(VDDerived);
-            if (Stmt* memsetCall = CheckAndBuildCallToMemset(
-                    BuildDeclRef(VDDerived),
-                    VDDerived->getInit()->IgnoreCasts()))
-              memsetCalls.push_back(memsetCall);
+            if (VDDerived->getInit() &&
+                !(VD->hasAttr<clang::CUDASharedAttr>() &&
+                  VD->getType()->isIncompleteArrayType())) {
+              if (Stmt* memsetCall = CheckAndBuildCallToMemset(
+                      BuildDeclRef(VDDerived),
+                      VDDerived->getInit()->IgnoreCasts()))
+                memsetCalls.push_back(memsetCall);
+            } else if (VD->hasAttr<clang::CUDASharedAttr>()) {
+              auto* VDForward = cast<clang::VarDecl>(decls.back());
+              HandleCUDASharedMemoryDecl(VD, VDForward, VDDerived, memsetCalls);
+            }
           }
         }
       } else if (auto* SAD = dyn_cast<StaticAssertDecl>(D)) {

lib/Differentiator/ReverseModeVisitorCUDA.cpp

@@ -0,0 +1,46 @@
+#include "ConstantFolder.h"
+#include "clad/Differentiator/ReverseModeVisitor.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/Attr.h"
+#include "clang/AST/Decl.h"
+
+using namespace clang;
+
+namespace clad {
+
+static void CloneCUDASharedAttr(const clang::VarDecl* OriginalVD,
+                                clang::VarDecl* VDClone) {
+  if (const auto* attr = OriginalVD->getAttr<clang::CUDASharedAttr>())
+    VDClone->addAttr(attr->clone(OriginalVD->getASTContext()));
+}
+
+void ReverseModeVisitor::HandleCUDASharedMemoryDecl(
+    const clang::VarDecl* VD, clang::VarDecl* VDForward,
+    clang::VarDecl* VDDerived,
+    llvm::SmallVectorImpl<clang::Stmt*>& memsetCalls) {
+
+  bool isDynamicSharedMem = VD->getType()->isIncompleteArrayType();
+
+  if (!isDynamicSharedMem) {
+    CloneCUDASharedAttr(VD, VDDerived);
+    VDDerived->setStorageClass(clang::SC_Static);
+
+    CloneCUDASharedAttr(VD, VDForward);
+    VDForward->setStorageClass(clang::SC_Static);
+
+    llvm::SmallVector<Expr*, 1> args = {BuildDeclRef(VDDerived)};
+    Stmt* initCall = GetCladZeroInit(args);
+    if (initCall)
+      memsetCalls.push_back(initCall);
+  } else {
+    CloneCUDASharedAttr(VD, VDForward);
+    Expr* derivedRef = BuildDeclRef(VDDerived);
+    Expr* zeroIdx =
+        ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, 0);
+    Expr* arraySub = BuildArraySubscript(derivedRef, {zeroIdx});
+    QualType elemType = VDDerived->getType()->getPointeeType();
+    Expr* assignZero = BuildOp(BO_Assign, arraySub, getZeroInit(elemType));
+    memsetCalls.push_back(assignZero);
+  }
+}
+} // namespace clad

lib/Differentiator/VisitorBase.cpp

@@ -180,9 +180,15 @@
   VarDecl* VisitorBase::BuildGlobalVarDecl(QualType Type,
                                            llvm::StringRef prefix, Expr* Init,
                                            bool DirectInit, TypeSourceInfo* TSI,
-                                           StorageClass SC) {
-    return BuildVarDecl(Type, CreateUniqueIdentifier(prefix),
-                        m_DerivativeFnScope, Init, DirectInit, TSI, SC);
+                                           StorageClass SC,
+                                           const VarDecl* OrigVD) {
+    VarDecl* VD = BuildVarDecl(Type, CreateUniqueIdentifier(prefix),
+                               m_DerivativeFnScope, Init, DirectInit, TSI, SC);
+    if (OrigVD) {
+      if (const auto* attr = OrigVD->getAttr<clang::CUDASharedAttr>())
+        VD->addAttr(attr->clone(OrigVD->getASTContext()));
+    }
+    return VD;
   }
 
   NamespaceDecl* VisitorBase::BuildNamespaceDecl(IdentifierInfo* II,
@@ -1086,4 +1092,4 @@
     }
   }
 
-} // end namespace clad
+  } // end namespace clad