From f01ec62e6660279364f20b05c558b599c067486a Mon Sep 17 00:00:00 2001 From: Vedant2005goyal Date: Fri, 12 Jun 2026 12:45:20 +0530 Subject: [PATCH] Add reverse-mode support for CUDA shared memory attributes and for the dynamic shared memory --- .../Differentiator/BuiltinDerivativesCUDA.cuh | 5 + .../clad/Differentiator/ReverseModeVisitor.h | 4 + include/clad/Differentiator/VisitorBase.h | 3 +- lib/Differentiator/CMakeLists.txt | 1 + lib/Differentiator/ReverseModeVisitor.cpp | 98 ++++++++++- lib/Differentiator/ReverseModeVisitorCUDA.cpp | 46 +++++ lib/Differentiator/VisitorBase.cpp | 14 +- test/CUDA/SharedMem.cu | 157 ++++++++++++++++++ 8 files changed, 316 insertions(+), 12 deletions(-) create mode 100644 lib/Differentiator/ReverseModeVisitorCUDA.cpp create mode 100644 test/CUDA/SharedMem.cu diff --git a/include/clad/Differentiator/BuiltinDerivativesCUDA.cuh b/include/clad/Differentiator/BuiltinDerivativesCUDA.cuh index ef5ec9df8..c9043228d 100644 --- a/include/clad/Differentiator/BuiltinDerivativesCUDA.cuh +++ b/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) { diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h index 32e34cacd..c17c835e4 100644 --- a/include/clad/Differentiator/ReverseModeVisitor.h +++ b/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& memsetCalls); }; } // end namespace clad diff --git a/include/clad/Differentiator/VisitorBase.h b/include/clad/Differentiator/VisitorBase.h index bdd2c29d3..519103f6c 100644 --- a/include/clad/Differentiator/VisitorBase.h +++ b/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. diff --git a/lib/Differentiator/CMakeLists.txt b/lib/Differentiator/CMakeLists.txt index bbb58a9d3..8a57e0704 100644 --- a/lib/Differentiator/CMakeLists.txt +++ b/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 diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index 3310e25ee..c887d8cff 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/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(E)) { + if (const auto* VD = dyn_cast(DRE->getDecl())) { + if (VD->hasAttr()) + return true; + } if (const auto* PVD = dyn_cast(DRE->getDecl())) { if (m_DiffReq->hasAttr()) // 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(E)) { const auto* base = dyn_cast(ASE->getBase()->IgnoreImpCasts()); + if (const auto* VD = dyn_cast(base->getDecl())) { + if (VD->hasAttr()) { + const auto* idx = ASE->getIdx(); + return !clad::utils::isInjective(idx, m_DiffReq.m_AnalysisDC); + } + } if (const auto* PVD = dyn_cast(base->getDecl())) { const auto* idx = ASE->getIdx(); if (m_DiffReq->hasAttr()) @@ -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()) { + 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(CE)) @@ -3052,9 +3072,20 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { dummyInit = m_Sema.ActOnInitList(noLoc, args, noLoc).get(); } } + bool isDynamicSharedMem = + VD->hasAttr() && VD->getType()->isIncompleteArrayType(); + + if (isDynamicSharedMem) { + QualType ElemTy = cast(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) && @@ -3062,8 +3093,20 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { 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(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()) + VDDerived->setInit(nullptr); } } @@ -3160,12 +3203,43 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { initDiff.getExpr(), VD->isDirectInit(), VDCloneTSI, SC); + if (isDynamicSharedMem && VDDerived) { + + llvm::SmallVector 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(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())) { @@ -3173,6 +3247,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { isConstructInit = true; } + if (!isDynamicSharedMem && VD->hasAttr()) + 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() && + VD->getType()->isIncompleteArrayType())) { + if (Stmt* memsetCall = CheckAndBuildCallToMemset( + BuildDeclRef(VDDerived), + VDDerived->getInit()->IgnoreCasts())) + memsetCalls.push_back(memsetCall); + } else if (VD->hasAttr()) { + auto* VDForward = cast(decls.back()); + HandleCUDASharedMemoryDecl(VD, VDForward, VDDerived, memsetCalls); + } } } } else if (auto* SAD = dyn_cast(D)) { diff --git a/lib/Differentiator/ReverseModeVisitorCUDA.cpp b/lib/Differentiator/ReverseModeVisitorCUDA.cpp new file mode 100644 index 000000000..b5af42168 --- /dev/null +++ b/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()) + VDClone->addAttr(attr->clone(OriginalVD->getASTContext())); +} + +void ReverseModeVisitor::HandleCUDASharedMemoryDecl( + const clang::VarDecl* VD, clang::VarDecl* VDForward, + clang::VarDecl* VDDerived, + llvm::SmallVectorImpl& 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 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 diff --git a/lib/Differentiator/VisitorBase.cpp b/lib/Differentiator/VisitorBase.cpp index aecef52c7..f8e416982 100644 --- a/lib/Differentiator/VisitorBase.cpp +++ b/lib/Differentiator/VisitorBase.cpp @@ -180,9 +180,15 @@ namespace clad { 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()) + VD->addAttr(attr->clone(OrigVD->getASTContext())); + } + return VD; } NamespaceDecl* VisitorBase::BuildNamespaceDecl(IdentifierInfo* II, @@ -1086,4 +1092,4 @@ namespace clad { } } -} // end namespace clad + } // end namespace clad \ No newline at end of file diff --git a/test/CUDA/SharedMem.cu b/test/CUDA/SharedMem.cu new file mode 100644 index 000000000..b5339c502 --- /dev/null +++ b/test/CUDA/SharedMem.cu @@ -0,0 +1,157 @@ +// RUN: %cladclang_cuda -I%S/../../include --cuda-path=%cudapath \ +// RUN: --cuda-gpu-arch=%cudaarch %cudaldflags -oReverseModeSharedMem.out \ +// RUN: -Xclang -verify %s 2>&1 | %filecheck %s +// +// RUN: %cudarun ./ReverseModeSharedMem.out | %filecheck_exec %s +// +// REQUIRES: cuda-runtime +// +// expected-no-diagnostics + +#include +#include "clad/Differentiator/Differentiator.h" +#include + +__global__ void func(int* vec_d, int val) { + __shared__ int sharedMem[1]; + sharedMem[0] = val * 3; + __syncthreads(); + vec_d[0] = sharedMem[0] * 5 + vec_d[0]; +} + +// CHECK: void func_grad(int *vec_d, int val, int *_d_vec_d, int *_d_val) { +// CHECK-NEXT: static int _d_sharedMem[1] __attribute__((shared)); +// CHECK-NEXT: clad::zero_init(_d_sharedMem); +// CHECK-NEXT: static int sharedMem[1] __attribute__((shared)); +// CHECK-NEXT: sharedMem[0] = val * 3; +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: vec_d[0] = sharedMem[0] * 5 + vec_d[0]; +// CHECK-NEXT: { +// CHECK-NEXT: int _r_d1 = _d_vec_d[0]; +// CHECK-NEXT: _d_vec_d[0] = 0; +// CHECK-NEXT: atomicAdd(&_d_sharedMem[0], _r_d1 * 5); +// CHECK-NEXT: atomicAdd(&_d_vec_d[0], _r_d1); +// CHECK-NEXT: } +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: { +// CHECK-NEXT: int _r_d0 = _d_sharedMem[0]; +// CHECK-NEXT: _d_sharedMem[0] = 0; +// CHECK-NEXT: atomicAdd(_d_val, _r_d0 * 3); +// CHECK-NEXT: } +// CHECK-NEXT: } + +__global__ void func1(int* vec1_d, int val1) { + extern __shared__ int sharedMem1[]; + sharedMem1[0] = val1 * 3; + __syncthreads(); + vec1_d[0] = sharedMem1[0] * 5 + vec1_d[0]; +} + +// CHECK: void func1_grad(int *vec1_d, int val1, int *_d_vec1_d, int *_d_val1) { +// CHECK-NEXT: int *_d_sharedMem1 = (int *)((char *)sharedMem1 + clad::get_dynamic_smem_size() / 2); +// CHECK-NEXT: _d_sharedMem1[0] = 0; +// CHECK-NEXT: extern int sharedMem1[] __attribute__((shared)); +// CHECK-NEXT: sharedMem1[0] = val1 * 3; +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: vec1_d[0] = sharedMem1[0] * 5 + vec1_d[0]; +// CHECK-NEXT: { +// CHECK-NEXT: int _r_d1 = _d_vec1_d[0]; +// CHECK-NEXT: _d_vec1_d[0] = 0; +// CHECK-NEXT: _d_sharedMem1[0] += _r_d1 * 5; +// CHECK-NEXT: atomicAdd(&_d_vec1_d[0], _r_d1); +// CHECK-NEXT: } +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: { +// CHECK-NEXT: int _r_d0 = _d_sharedMem1[0]; +// CHECK-NEXT: _d_sharedMem1[0] = 0; +// CHECK-NEXT: atomicAdd(_d_val1, _r_d0 * 3); +// CHECK-NEXT: } +// CHECK-NEXT: } + +__global__ void scalar_func(int* vec2_d, int val2) { + __shared__ int x; + x = val2 * 3; + __syncthreads(); + vec2_d[0] = x * 5; +} + +// CHECK: void scalar_func_grad(int *vec2_d, int val2, int *_d_vec2_d, int *_d_val2) { +// CHECK-NEXT: static int _d_x __attribute__((shared)); +// CHECK-NEXT: clad::zero_init(_d_x); +// CHECK-NEXT: static int x __attribute__((shared)); +// CHECK-NEXT: x = val2 * 3; +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: vec2_d[0] = x * 5; +// CHECK-NEXT: { +// CHECK-NEXT: int _r_d0 = _d_vec2_d[0]; +// CHECK-NEXT: _d_vec2_d[0] = 0; +// CHECK-NEXT: atomicAdd(&_d_x, _r_d0 * 5); +// CHECK-NEXT: } +// CHECK-NEXT: __syncthreads(); +// CHECK-NEXT: { +// CHECK-NEXT: atomicAdd(_d_val2, _d_x * 3); +// CHECK-NEXT: _d_x = 0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +int main() { + + int *d_vec, *d_vec_adj, *d_val_adj; + cudaMalloc(&d_vec, sizeof(int)); + cudaMalloc(&d_vec_adj, sizeof(int)); + cudaMalloc(&d_val_adj, sizeof(int)); + + int seed = 1, seed1 = 0, grad, grad1; + + cudaMemcpy(d_vec_adj, &seed, sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_val_adj, &seed1, sizeof(int), cudaMemcpyHostToDevice); + + auto dfunc = clad::gradient(func); + dfunc.execute_kernel(dim3(1), dim3(1), d_vec, 90, d_vec_adj, d_val_adj); + cudaDeviceSynchronize(); + + cudaMemcpy(&grad, d_val_adj, sizeof(int), cudaMemcpyDeviceToHost); + cudaMemcpy(&grad1, d_vec_adj, sizeof(int), cudaMemcpyDeviceToHost); + + std::cout << "Grad dvec/dval: " << grad << std::endl; + // CHECK-EXEC: Grad dvec/dval: 15 + std::cout << "Grad dvec/dvec_d[0]: " << grad1 << std::endl; + // CHECK-EXEC-NEXT: Grad dvec/dvec_d[0]: 1 + + cudaMemcpy(d_vec_adj, &seed, sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_val_adj, &seed1, sizeof(int), cudaMemcpyHostToDevice); + + auto dfunc1 = clad::gradient(func1); + auto* kernel_ptr = dfunc1.getFunctionPtr(); + kernel_ptr<<>>(d_vec, 90, d_vec_adj, d_val_adj); + cudaDeviceSynchronize(); + + cudaMemcpy(&grad, d_val_adj, sizeof(int), cudaMemcpyDeviceToHost); + cudaMemcpy(&grad1, d_vec_adj, sizeof(int), cudaMemcpyDeviceToHost); + + std::cout << "Grad dvec/dval: " << grad << std::endl; + // CHECK-EXEC-NEXT: Grad dvec/dval: 15 + std::cout << "Grad dvec/dvec1_d[0]: " << grad1 << std::endl; + // CHECK-EXEC-NEXT: Grad dvec/dvec1_d[0]: 1 + + cudaMemcpy(d_vec_adj, &seed, sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_val_adj, &seed1, sizeof(int), cudaMemcpyHostToDevice); + + auto d_scalar = clad::gradient(scalar_func); + d_scalar.execute_kernel(dim3(1), dim3(1), d_vec, 9, d_vec_adj, d_val_adj); + cudaDeviceSynchronize(); + + cudaMemcpy(&grad, d_val_adj, sizeof(int), cudaMemcpyDeviceToHost); + cudaMemcpy(&grad1, d_vec_adj, sizeof(int), cudaMemcpyDeviceToHost); + + std::cout << "Grad dvec/dval: " << grad << std::endl; + // CHECK-EXEC-NEXT: Grad dvec/dval: 15 + std::cout << "Grad dvec/dvec2_d[0]: " << grad1 << std::endl; + // CHECK-EXEC-NEXT: Grad dvec/dvec2_d[0]: 0 + + cudaFree(d_vec); + cudaFree(d_vec_adj); + cudaFree(d_val_adj); + + return 0; +}