diff --git a/include/clad/Differentiator/CladUtils.h b/include/clad/Differentiator/CladUtils.h index 64e07269f..18d1876fa 100644 --- a/include/clad/Differentiator/CladUtils.h +++ b/include/clad/Differentiator/CladUtils.h @@ -408,6 +408,10 @@ namespace clad { /// create modifiable adjoints. clang::QualType replaceStdInitListWithCladArray(clang::Sema& S, clang::QualType origTy); + /// Currently is only used for CUDA in the reverse mode. Determines whether + /// an expression, most likely an index, is injective, meaning no two + /// threads have the same value. + bool isInjective(const clang::Expr* E, clang::ASTContext& ctx); } // namespace utils } // namespace clad diff --git a/lib/Differentiator/CladUtils.cpp b/lib/Differentiator/CladUtils.cpp index 6e585acc5..6069dc060 100644 --- a/lib/Differentiator/CladUtils.cpp +++ b/lib/Differentiator/CladUtils.cpp @@ -8,6 +8,7 @@ #include "clang/AST/DeclCXX.h" #include "clang/AST/Expr.h" #include "clang/AST/ExprCXX.h" +#include "clang/AST/OperationKinds.h" #include "clang/AST/RecursiveASTVisitor.h" #include "clang/AST/Type.h" #include "clang/Basic/Builtins.h" @@ -17,6 +18,9 @@ #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" +#include +#include + using namespace clang; namespace clad { namespace utils { @@ -1144,5 +1148,195 @@ namespace clad { return S.getASTContext().getLValueReferenceType(T); return T; } + + static bool isInjectiveE(const clang::Expr* E) { + class InjectiveCheckerExpr + : public clang::RecursiveASTVisitor { + struct IdxNode { + struct ExprOrBinOp { + const clang::DeclRefExpr* m_E = nullptr; + clang::BinaryOperatorKind m_Opcode; + enum class Kind { Expr, Opcode } m_Kind; + + ExprOrBinOp(const clang::DeclRefExpr* E) + : m_E(E), m_Kind(Kind::Expr) {} + ExprOrBinOp(clang::BinaryOperatorKind op) + : m_Opcode(op), m_Kind(Kind::Opcode) {} + + [[nodiscard]] bool isExpr() const { return m_Kind == Kind::Expr; } + [[nodiscard]] bool isOpcode() const { + return m_Kind == Kind::Opcode; + } + }; + + ExprOrBinOp Node; + std::unique_ptr left; + std::unique_ptr right; + + IdxNode(ExprOrBinOp N) : Node(N) {} + }; + std::unique_ptr m_Root; + IdxNode* m_ParentNode = nullptr; + + enum class side { left, right } m_Side; + + public: + InjectiveCheckerExpr() = default; + /// This function recursively checks whether a given pattern matches the + /// previously computed graph. It uses a fairly standard graph + /// comparison algorithm. + bool comparePatternToTree(const IdxNode* current, + const std::vector& patternIdx, + size_t i = 1) { + // If current is not initialized and child is empty or does not exist, + // we have a match. + if (!current && (i >= patternIdx.size() || patternIdx[i].empty())) + return true; + + if (!current || i >= patternIdx.size() || patternIdx[i].empty()) + return false; + + const std::string& expected = patternIdx[i]; + + if (current->Node.isOpcode()) { + std::string actualOp = + clang::BinaryOperator::getOpcodeStr(current->Node.m_Opcode) + .str(); + if (actualOp != expected) + return false; + } else if (current->Node.isExpr()) { + std::string actualName = + current->Node.m_E->getNameInfo().getAsString(); + if (actualName != expected) + return false; + } + + bool sameOrder = + comparePatternToTree(current->left.get(), patternIdx, 2 * i) && + comparePatternToTree(current->right.get(), patternIdx, 2 * i + 1); + + if (sameOrder) + return true; + // Here we account for a sub-tree rotation wrt to the current node. If + // there is no match at this point, we compare a pattern to a graph + // with the rotation. + return comparePatternToTree(current->left.get(), patternIdx, + 2 * i + 1) && + comparePatternToTree(current->right.get(), patternIdx, 2 * i); + } + + bool isInjectiveIdx(const clang::Expr* E) { + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) + TraverseStmt(const_cast(E)); + std::vector pattern = {"", "+", "threadIdx", "*", + "", "", "blockIdx", "blockDim"}; + return comparePatternToTree(m_Root.get(), pattern); + } + + bool TraverseBinaryOperator(clang::BinaryOperator* BinOp) { + const auto opCode = BinOp->getOpcode(); + if (opCode == BO_Add || opCode == BO_Mul) { + std::unique_ptr curr = std::make_unique(opCode); + IdxNode* currPtr = nullptr; + + if (!m_Root) { + m_Root = std::move(curr); + currPtr = m_Root.get(); + } else { + currPtr = curr.get(); + + if (m_Side == side::left) + m_ParentNode->left = std::move(curr); + + if (m_Side == side::right) + m_ParentNode->right = std::move(curr); + } + + Expr* L = BinOp->getLHS(); + Expr* R = BinOp->getRHS(); + + m_ParentNode = currPtr; + m_Side = side::left; + TraverseStmt(L); + m_ParentNode = currPtr; + + m_Side = side::right; + TraverseStmt(R); + } + return true; + } + + bool TraverseDeclRefExpr(clang::DeclRefExpr* DRE) { + std::unique_ptr curr = std::make_unique(DRE); + + if (m_ParentNode) { + if (m_Side == side::left) + m_ParentNode->left = std::move(curr); + + if (m_Side == side::right) + m_ParentNode->right = std::move(curr); + } + return true; + } + + } checker; + return checker.isInjectiveIdx(E); + } + + bool isInjective(const clang::Expr* E, clang::ASTContext& ctx) { + class InjectiveChecker + : public clang::RecursiveASTVisitor { + clang::ASTContext& m_Context; + + public: + InjectiveChecker(clang::ASTContext& Context) : m_Context(Context) {}; + + bool isInjectiveIdx(const clang::Expr* E) { + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) + return TraverseStmt(const_cast(E)); + } + + bool TraverseBinaryOperator(clang::BinaryOperator* BinOp) { + const auto opCode = BinOp->getOpcode(); + Expr* L = BinOp->getLHS(); + Expr* R = BinOp->getRHS(); + + if (opCode == BO_Add || opCode == BO_Mul) { + Expr::EvalResult dummy; + + bool isConstL = + clad_compat::Expr_EvaluateAsConstantExpr(L, dummy, m_Context); + bool isConstR = + clad_compat::Expr_EvaluateAsConstantExpr(R, dummy, m_Context); + + if (isConstL && isConstR) + return false; + + if (!isConstL && isConstR) + return TraverseStmt(L); + + if (!isConstL && !isConstR) + return isInjectiveE(BinOp); + + if (!isConstR) + return TraverseStmt(R); + } + return false; + } + + bool TraverseDeclRefExpr(clang::DeclRefExpr* DRE) { + if (auto* VD = dyn_cast(DRE->getDecl())) { + if (auto* init = VD->getInit()) + return isInjectiveE(init->IgnoreImpCasts()); + } + return false; + } + + bool TraverseIntegerLiteral(clang::IntegerLiteral* IL) { return false; } + + } checker(ctx); + + return checker.isInjectiveIdx(E); + } } // namespace utils } // namespace clad diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index aff8c939d..7ce064fcd 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/lib/Differentiator/ReverseModeVisitor.cpp @@ -52,13 +52,15 @@ #include #include #include +#include #include +#include +#include +#include #include "clad/Differentiator/CladUtils.h" #include "clad/Differentiator/Compatibility.h" -using namespace clang; - namespace clad { Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { @@ -126,27 +128,42 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { bool ReverseModeVisitor::shouldUseCudaAtomicOps(const Expr* E) { if (!m_Context.getLangOpts().CUDA) return false; - - if (!isa(E)) - return false; - - const auto* DRE = cast(E); - - if (const auto* PVD = dyn_cast(DRE->getDecl())) { - if (m_DiffReq->hasAttr()) - // Check whether this param is in the global memory of the GPU - return m_DiffReq.HasIndependentParameter(PVD); - if (m_DiffReq->hasAttr()) { - for (auto index : m_DiffReq.CUDAGlobalArgsIndexes) { - const auto* PVDOrig = m_DiffReq->getParamDecl(index); - if ("_d_" + PVDOrig->getNameAsString() == PVD->getNameAsString() && - (utils::isArrayOrPointerType(PVDOrig->getType()) || - PVDOrig->getType()->isReferenceType())) - return true; + if (const auto* DRE = dyn_cast(E)) { + if (const auto* PVD = dyn_cast(DRE->getDecl())) { + if (m_DiffReq->hasAttr()) + // Check whether this param is in the global memory of the GPU + return m_DiffReq.HasIndependentParameter(PVD); + if (m_DiffReq->hasAttr()) { + for (auto index : m_DiffReq.CUDAGlobalArgsIndexes) { + const auto* PVDOrig = m_DiffReq->getParamDecl(index); + if ("_d_" + PVDOrig->getNameAsString() == PVD->getNameAsString() && + (utils::isArrayOrPointerType(PVDOrig->getType()) || + PVDOrig->getType()->isReferenceType())) + return true; + } + } + } + } else if (const auto* ASE = dyn_cast(E)) { + const auto* base = + dyn_cast(ASE->getBase()->IgnoreImpCasts()); + if (const auto* PVD = dyn_cast(base->getDecl())) { + const auto* idx = ASE->getIdx(); + if (m_DiffReq->hasAttr()) + // Check whether this param is in the global memory of the GPU and + // if index is injective. + return m_DiffReq.HasIndependentParameter(PVD) && + !clad::utils::isInjective(idx, m_Context); + if (m_DiffReq->hasAttr()) { + for (auto index : m_DiffReq.CUDAGlobalArgsIndexes) { + const auto* PVDOrig = m_DiffReq->getParamDecl(index); + if (PVDOrig->getNameAsString() == PVD->getNameAsString() && + (utils::isArrayOrPointerType(PVDOrig->getType()) || + PVDOrig->getType()->isReferenceType())) + return !clad::utils::isInjective(idx, m_Context); + } } } } - return false; } @@ -1376,7 +1393,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { // Create the (target += dfdx) statement. if (dfdx()) { Expr* add_assign = nullptr; - if (shouldUseCudaAtomicOps(target)) + if (shouldUseCudaAtomicOps(ASE)) add_assign = BuildCallToCudaAtomicAdd(result, dfdx()); else add_assign = BuildOp(BO_AddAssign, result, dfdx()); diff --git a/test/CUDA/GradientKernels.cu b/test/CUDA/GradientKernels.cu index 0d0bd614c..c81847fd3 100644 --- a/test/CUDA/GradientKernels.cu +++ b/test/CUDA/GradientKernels.cu @@ -77,7 +77,7 @@ __global__ void add_kernel_3(int *out, int *in) { //CHECK-NEXT: { //CHECK-NEXT: out[index0] = _t0; //CHECK-NEXT: int _r_d0 = _d_out[index0]; -//CHECK-NEXT: atomicAdd(&_d_in[index0], _r_d0); +//CHECK-NEXT: _d_in[index0] += _r_d0; //CHECK-NEXT: } //CHECK-NEXT:} @@ -347,7 +347,7 @@ __global__ void dup_kernel_with_device_call_2(double *out, const double *in, dou //CHECK-NEXT: int _d_index = 0; //CHECK-NEXT: int index0 = threadIdx.x + blockIdx.x * blockDim.x; //CHECK-NEXT: { -//CHECK-NEXT: atomicAdd(&_d_in[index0], _d_y); +//CHECK-NEXT: _d_in[index0] += _d_y; //CHECK-NEXT: *_d_val += _d_y; //CHECK-NEXT: } //CHECK-NEXT:} @@ -382,7 +382,7 @@ __global__ void kernel_with_device_call_3(double *out, double *in, double *val) //CHECK-NEXT: int _d_index = 0; //CHECK-NEXT: int index0 = threadIdx.x + blockIdx.x * blockDim.x; //CHECK-NEXT: { -//CHECK-NEXT: atomicAdd(&_d_in[index0], _d_y); +//CHECK-NEXT: _d_in[index0] += _d_y; //CHECK-NEXT: atomicAdd(_d_val, _d_y); //CHECK-NEXT: } //CHECK-NEXT:} @@ -418,7 +418,7 @@ __global__ void kernel_with_nested_device_call(double *out, double *in, double v //CHECK-NEXT: int _d_index = 0; //CHECK-NEXT: int index0 = threadIdx.x + blockIdx.x * blockDim.x; //CHECK-NEXT: { -//CHECK-NEXT: atomicAdd(&_d_in[index0], _d_y); +//CHECK-NEXT: _d_in[index0] += _d_y; //CHECK-NEXT: *_d_val += _d_y; //CHECK-NEXT: } //CHECK-NEXT:} @@ -689,6 +689,154 @@ void launch_add_kernel_4(int *out, int *in, const int N) { //CHECK-NEXT: cudaFree(_d_out_dev); //CHECK-NEXT:} +__global__ void indices_perm(int *out, int *in) { + int index1 = threadIdx.x + blockIdx.x * blockDim.x; + int index2 = threadIdx.x + blockDim.x * blockIdx.x; + int index3 = blockIdx.x * blockDim.x + threadIdx.x; + int index4 = blockDim.x * blockIdx.x + threadIdx.x; + out[index1] += in[index1]; + out[index2] += in[index2]; + out[index3] += in[index3]; + out[index4] += in[index4]; +} + +// CHECK: void indices_perm_grad(int *out, int *in, int *_d_out, int *_d_in) { +// CHECK-NEXT: int _d_index1 = 0; +// CHECK-NEXT: int index1 = threadIdx.x + blockIdx.x * blockDim.x; +// CHECK-NEXT: int _d_index2 = 0; +// CHECK-NEXT: int index2 = threadIdx.x + blockDim.x * blockIdx.x; +// CHECK-NEXT: int _d_index3 = 0; +// CHECK-NEXT: int index3 = blockIdx.x * blockDim.x + threadIdx.x; +// CHECK-NEXT: int _d_index4 = 0; +// CHECK-NEXT: int index4 = blockDim.x * blockIdx.x + threadIdx.x; +// CHECK-NEXT: int _t0 = out[index1]; +// CHECK-NEXT: out[index1] += in[index1]; +// CHECK-NEXT: int _t1 = out[index2]; +// CHECK-NEXT: out[index2] += in[index2]; +// CHECK-NEXT: int _t2 = out[index3]; +// CHECK-NEXT: out[index3] += in[index3]; +// CHECK-NEXT: int _t3 = out[index4]; +// CHECK-NEXT: out[index4] += in[index4]; +// CHECK-NEXT: { +// CHECK-NEXT: out[index4] = _t3; +// CHECK-NEXT: int _r_d3 = _d_out[index4]; +// CHECK-NEXT: _d_in[index4] += _r_d3; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index3] = _t2; +// CHECK-NEXT: int _r_d2 = _d_out[index3]; +// CHECK-NEXT: _d_in[index3] += _r_d2; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index2] = _t1; +// CHECK-NEXT: int _r_d1 = _d_out[index2]; +// CHECK-NEXT: _d_in[index2] += _r_d1; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index1] = _t0; +// CHECK-NEXT: int _r_d0 = _d_out[index1]; +// CHECK-NEXT: _d_in[index1] += _r_d0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +__global__ void indices_lin_comb(int *out, int *in) { + int index = threadIdx.x + blockIdx.x * blockDim.x; + + out[index] += in[2*index]; + out[index] += in[1+index]; + out[index] += in[threadIdx.x + blockIdx.x * blockDim.x]; + out[index] += in[2*(threadIdx.x + blockIdx.x * blockDim.x) + 1]; + out[index] += in[1+1]; + out[index] += in[index/2]; + +} + +// CHECK: void indices_lin_comb_grad(int *out, int *in, int *_d_out, int *_d_in) { +// CHECK-NEXT: int _d_index = 0; +// CHECK-NEXT: int index0 = threadIdx.x + blockIdx.x * blockDim.x; +// CHECK-NEXT: int _t0 = out[index0]; +// CHECK-NEXT: out[index0] += in[2 * index0]; +// CHECK-NEXT: int _t1 = out[index0]; +// CHECK-NEXT: out[index0] += in[1 + index0]; +// CHECK-NEXT: int _t2 = out[index0]; +// CHECK-NEXT: out[index0] += in[threadIdx.x + blockIdx.x * blockDim.x]; +// CHECK-NEXT: int _t3 = out[index0]; +// CHECK-NEXT: unsigned int _t4 = (threadIdx.x + blockIdx.x * blockDim.x); +// CHECK-NEXT: out[index0] += in[2 * _t4 + 1]; +// CHECK-NEXT: int _t5 = out[index0]; +// CHECK-NEXT: out[index0] += in[1 + 1]; +// CHECK-NEXT: int _t6 = out[index0]; +// CHECK-NEXT: out[index0] += in[index0 / 2]; +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t6; +// CHECK-NEXT: int _r_d5 = _d_out[index0]; +// CHECK-NEXT: atomicAdd(&_d_in[index0 / 2], _r_d5); +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t5; +// CHECK-NEXT: int _r_d4 = _d_out[index0]; +// CHECK-NEXT: atomicAdd(&_d_in[1 + 1], _r_d4); +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t3; +// CHECK-NEXT: int _r_d3 = _d_out[index0]; +// CHECK-NEXT: _d_in[2 * _t4 + 1] += _r_d3; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t2; +// CHECK-NEXT: int _r_d2 = _d_out[index0]; +// CHECK-NEXT: _d_in[threadIdx.x + blockIdx.x * blockDim.x] += _r_d2; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t1; +// CHECK-NEXT: int _r_d1 = _d_out[index0]; +// CHECK-NEXT: _d_in[1 + index0] += _r_d1; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: out[index0] = _t0; +// CHECK-NEXT: int _r_d0 = _d_out[index0]; +// CHECK-NEXT: _d_in[2 * index0] += _r_d0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +__device__ void device_injective_index(int *a) { + int index1 = threadIdx.x + blockIdx.x * blockDim.x; + int index2 = threadIdx.x; + a[index1] += a[index1]; + a[index2] += a[index2]; +} + +__global__ void kernel_device_injective(int *a) { + device_injective_index(a); +} + +// CHECK: __attribute__((device)) void device_injective_index_pullback_0(int *a, int *_d_a) { +// CHECK-NEXT: int _d_index1 = 0; +// CHECK-NEXT: int index1 = threadIdx.x + blockIdx.x * blockDim.x; +// CHECK-NEXT: int _d_index2 = 0; +// CHECK-NEXT: int index2 = threadIdx.x; +// CHECK-NEXT: int _t0 = a[index1]; +// CHECK-NEXT: a[index1] += a[index1]; +// CHECK-NEXT: int _t1 = a[index2]; +// CHECK-NEXT: a[index2] += a[index2]; +// CHECK-NEXT: { +// CHECK-NEXT: a[index2] = _t1; +// CHECK-NEXT: int _r_d1 = _d_a[index2]; +// CHECK-NEXT: atomicAdd(&_d_a[index2], _r_d1); +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: a[index1] = _t0; +// CHECK-NEXT: int _r_d0 = _d_a[index1]; +// CHECK-NEXT: _d_a[index1] += _r_d0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +// CHECK: void kernel_device_injective_grad(int *a, int *_d_a) { +// CHECK-NEXT: device_injective_index(a); +// CHECK-NEXT: device_injective_index_pullback_0(a, _d_a); +// CHECK-NEXT: } + + #define TEST(F, grid, block, shared_mem, use_stream, x, dx, N) \ { \ int *fives = (int*)malloc(N * sizeof(int)); \ @@ -963,6 +1111,17 @@ int main(void) { launch_kernel_4_test.execute(zeros_int, fives_int, 10, out_res, in_res); printf("%d, %d, %d\n", in_res[0], in_res[1], in_res[2]); // CHECK-EXEC: 5, 5, 5 + TEST_2(indices_perm, dim3(1), dim3(5, 1, 1), 0, false, "out, in", dummy_out, dummy_in, d_out, d_in, 5); // CHECK-EXEC: 20, 20, 20, 20, 20 + TEST_2(indices_lin_comb, dim3(1), dim3(5, 1, 1), 0, false, "out, in", dummy_out, dummy_in, d_out, d_in, 5); // CHECK-EXEC: 20, 25, 45, 15, 15 + + int *n, *d_n; + cudaMalloc(&n, sizeof(int)); + cudaMalloc(&d_n, sizeof(int)); + + TEST(kernel_device_injective, dim3(1), dim3(1), 0, false, n, d_n, 1); // CHECK-EXEC: 4 + + cudaFree(n); + cudaFree(d_n); free(res); free(fives); free(zeros);