Skip to content

Commit c4dbec5

Browse files
author
ovdiiuv
committed
Rework isInjectiveE not to use ASTMatchers
1 parent 0f3a496 commit c4dbec5

2 files changed

Lines changed: 241 additions & 50 deletions

File tree

‎lib/Differentiator/ReverseModeVisitor.cpp‎

Lines changed: 128 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,15 @@
5151
#include <algorithm>
5252
#include <cstddef>
5353
#include <iterator>
54+
#include <memory>
5455
#include <numeric>
56+
#include <string>
57+
#include <utility>
58+
#include <vector>
5559

5660
#include "clad/Differentiator/CladUtils.h"
5761
#include "clad/Differentiator/Compatibility.h"
5862

59-
#include "clang/ASTMatchers/ASTMatchFinder.h"
60-
#include "clang/ASTMatchers/ASTMatchers.h"
61-
62-
using namespace clang::ast_matchers;
63-
6463
namespace clad {
6564

6665
Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
@@ -136,51 +135,130 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
136135
return CladTapeResult{*this, PushExpr, PopExpr, TapeRef};
137136
}
138137

139-
static bool isInjectiveE(const clang::Expr* E, clang::ASTContext& ctx) {
140-
StatementMatcher matcher =
141-
binaryOperator(
142-
hasOperatorName("+"),
143-
hasLHS(expr(
144-
has(callExpr(hasDescendant(memberExpr(
145-
hasObjectExpression(opaqueValueExpr(hasSourceExpression(
146-
declRefExpr(to(varDecl(hasName("threadIdx"))))))))))),
147-
unless(binaryOperator()))),
148-
hasRHS(binaryOperator(
149-
hasOperatorName("*"),
150-
has(expr(
151-
has(callExpr(hasDescendant(memberExpr(hasObjectExpression(
152-
opaqueValueExpr(hasSourceExpression(declRefExpr(
153-
to(varDecl(hasName("blockIdx"))))))))))),
154-
unless(binaryOperator()))),
155-
has(expr(
156-
has(callExpr(hasDescendant(memberExpr(hasObjectExpression(
157-
opaqueValueExpr(hasSourceExpression(declRefExpr(
158-
to(varDecl(hasName("blockDim"))))))))))),
159-
unless(binaryOperator()))))))
160-
.bind("targetVar");
161-
162-
class InjectivePatternMatchCallback : public MatchFinder::MatchCallback {
163-
const clang::Expr* targetExpr;
138+
static bool isInjectiveE(const clang::Expr* E) {
139+
class InjectiveCheckerExpr
140+
: public clang::RecursiveASTVisitor<InjectiveCheckerExpr> {
141+
struct IdxNode {
142+
struct ExprOrBinOp {
143+
const clang::DeclRefExpr* m_E = nullptr;
144+
clang::BinaryOperatorKind m_Opcode;
145+
enum class Kind { Expr, Opcode } m_Kind;
146+
147+
ExprOrBinOp(const clang::DeclRefExpr* E)
148+
: m_E(E), m_Kind(Kind::Expr) {}
149+
ExprOrBinOp(clang::BinaryOperatorKind op)
150+
: m_Opcode(op), m_Kind(Kind::Opcode) {}
151+
152+
[[nodiscard]] bool isExpr() const { return m_Kind == Kind::Expr; }
153+
[[nodiscard]] bool isOpcode() const { return m_Kind == Kind::Opcode; }
154+
};
155+
156+
ExprOrBinOp Node;
157+
std::unique_ptr<IdxNode> left;
158+
std::unique_ptr<IdxNode> right;
159+
160+
IdxNode(ExprOrBinOp N) : Node(N) {}
161+
};
162+
std::unique_ptr<IdxNode> m_Root;
163+
IdxNode* m_ParentNode = nullptr;
164+
165+
enum class side { left, right } m_Side;
164166

165167
public:
166-
InjectivePatternMatchCallback(const clang::Expr* E) : targetExpr(E) {}
167-
bool matched = false;
168-
virtual void run(const MatchFinder::MatchResult& Result) override {
169-
if (const auto* expr =
170-
Result.Nodes.getNodeAs<clang::BinaryOperator>("targetVar")) {
171-
if (targetExpr == expr)
172-
matched = true;
168+
InjectiveCheckerExpr() = default;
169+
170+
bool comparePatternToTree(const IdxNode* current,
171+
const std::vector<std::string>& patternIdx,
172+
size_t i = 1) {
173+
174+
if (!current && (i >= patternIdx.size() || patternIdx[i].empty()))
175+
return true;
176+
177+
if (!current || i >= patternIdx.size() || patternIdx[i].empty())
178+
return false;
179+
180+
const std::string& expected = patternIdx[i];
181+
182+
if (current->Node.isOpcode()) {
183+
std::string actualOp =
184+
clang::BinaryOperator::getOpcodeStr(current->Node.m_Opcode).str();
185+
if (actualOp != expected)
186+
return false;
187+
} else if (current->Node.isExpr()) {
188+
std::string actualName =
189+
current->Node.m_E->getNameInfo().getAsString();
190+
if (actualName != expected)
191+
return false;
173192
}
193+
194+
bool sameOrder =
195+
comparePatternToTree(current->left.get(), patternIdx, 2 * i) &&
196+
comparePatternToTree(current->right.get(), patternIdx, 2 * i + 1);
197+
198+
if (sameOrder)
199+
return true;
200+
201+
return comparePatternToTree(current->left.get(), patternIdx,
202+
2 * i + 1) &&
203+
comparePatternToTree(current->right.get(), patternIdx, 2 * i);
204+
}
205+
206+
bool isInjectiveIdx(const clang::Expr* E) {
207+
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
208+
TraverseStmt(const_cast<clang::Expr*>(E));
209+
std::vector<std::string> pattern = {"", "+", "threadIdx", "*",
210+
"", "", "blockIdx", "blockDim"};
211+
return comparePatternToTree(m_Root.get(), pattern);
212+
}
213+
214+
bool TraverseBinaryOperator(clang::BinaryOperator* BinOp) {
215+
const auto opCode = BinOp->getOpcode();
216+
if (opCode == BO_Add || opCode == BO_Mul) {
217+
std::unique_ptr<IdxNode> curr = std::make_unique<IdxNode>(opCode);
218+
IdxNode* currPtr = nullptr;
219+
220+
if (!m_Root) {
221+
m_Root = std::move(curr);
222+
currPtr = m_Root.get();
223+
} else {
224+
currPtr = curr.get();
225+
226+
if (m_Side == side::left)
227+
m_ParentNode->left = std::move(curr);
228+
229+
if (m_Side == side::right)
230+
m_ParentNode->right = std::move(curr);
231+
}
232+
233+
Expr* L = BinOp->getLHS();
234+
Expr* R = BinOp->getRHS();
235+
236+
m_ParentNode = currPtr;
237+
m_Side = side::left;
238+
TraverseStmt(L);
239+
m_ParentNode = currPtr;
240+
241+
m_Side = side::right;
242+
TraverseStmt(R);
243+
}
244+
return true;
174245
}
175-
bool hasMatched() const { return matched; }
176-
};
177246

178-
MatchFinder Finder;
179-
InjectivePatternMatchCallback Callback(E);
180-
Finder.addMatcher(matcher, &Callback);
181-
Finder.matchAST(ctx);
247+
bool TraverseDeclRefExpr(clang::DeclRefExpr* DRE) {
248+
std::unique_ptr<IdxNode> curr = std::make_unique<IdxNode>(DRE);
249+
250+
if (m_ParentNode) {
251+
if (m_Side == side::left)
252+
m_ParentNode->left = std::move(curr);
182253

183-
return Callback.hasMatched();
254+
if (m_Side == side::right)
255+
m_ParentNode->right = std::move(curr);
256+
}
257+
return true;
258+
}
259+
260+
} checker;
261+
return checker.isInjectiveIdx(E);
184262
}
185263

186264
static bool isInjective(const clang::Expr* E, clang::ASTContext& ctx) {
@@ -189,7 +267,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
189267
clang::ASTContext& m_Context;
190268

191269
public:
192-
InjectiveChecker(clang::ASTContext& Context) : m_Context(Context){};
270+
InjectiveChecker(clang::ASTContext& Context) : m_Context(Context) {};
193271

194272
bool isInjectiveIdx(const clang::Expr* E) {
195273
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
@@ -218,15 +296,15 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
218296
if (isConstL && !isConstR)
219297
return TraverseStmt(R);
220298

221-
return isInjectiveE(BinOp, m_Context);
299+
return isInjectiveE(BinOp);
222300
}
223301
return false;
224302
}
225303

226304
bool TraverseDeclRefExpr(clang::DeclRefExpr* DRE) {
227305
if (auto* VD = dyn_cast<VarDecl>(DRE->getDecl())) {
228306
if (auto* init = VD->getInit())
229-
return isInjectiveE(init->IgnoreImpCasts(), m_Context);
307+
return isInjectiveE(init->IgnoreImpCasts());
230308
}
231309
return false;
232310
}
@@ -258,8 +336,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
258336
}
259337
} else if (const auto* ASE = dyn_cast<ArraySubscriptExpr>(E)) {
260338
if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>()) {
261-
auto* idx = ASE->getIdx();
262-
auto* base = dyn_cast<DeclRefExpr>(ASE->getBase()->IgnoreImpCasts());
339+
const auto* idx = ASE->getIdx();
340+
const auto* base =
341+
dyn_cast<DeclRefExpr>(ASE->getBase()->IgnoreImpCasts());
263342
if (const auto* PVD = dyn_cast<ParmVarDecl>(base->getDecl())) {
264343
if (m_DiffReq->hasAttr<clang::CUDAGlobalAttr>()) {
265344
// Check whether this param is in the global memory of the GPU and
@@ -268,7 +347,6 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
268347
!isInjective(idx, m_Context);
269348
}
270349
}
271-
return true;
272350
}
273351
}
274352
return false;

‎test/CUDA/GradientKernels.cu‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -689,6 +689,116 @@ void launch_add_kernel_4(int *out, int *in, const int N) {
689689
//CHECK-NEXT: cudaFree(_d_out_dev);
690690
//CHECK-NEXT:}
691691

692+
__global__ void indices_perm(int *out, int *in) {
693+
int index1 = threadIdx.x + blockIdx.x * blockDim.x;
694+
int index2 = threadIdx.x + blockDim.x * blockIdx.x;
695+
int index3 = blockIdx.x * blockDim.x + threadIdx.x;
696+
int index4 = blockDim.x * blockIdx.x + threadIdx.x;
697+
out[index1] += in[index1];
698+
out[index2] += in[index2];
699+
out[index3] += in[index3];
700+
out[index4] += in[index4];
701+
}
702+
703+
// CHECK: void indices_perm_grad(int *out, int *in, int *_d_out, int *_d_in) {
704+
// CHECK-NEXT: int _d_index1 = 0;
705+
// CHECK-NEXT: int index1 = threadIdx.x + blockIdx.x * blockDim.x;
706+
// CHECK-NEXT: int _d_index2 = 0;
707+
// CHECK-NEXT: int index2 = threadIdx.x + blockDim.x * blockIdx.x;
708+
// CHECK-NEXT: int _d_index3 = 0;
709+
// CHECK-NEXT: int index3 = blockIdx.x * blockDim.x + threadIdx.x;
710+
// CHECK-NEXT: int _d_index4 = 0;
711+
// CHECK-NEXT: int index4 = blockDim.x * blockIdx.x + threadIdx.x;
712+
// CHECK-NEXT: int _t0 = out[index1];
713+
// CHECK-NEXT: out[index1] += in[index1];
714+
// CHECK-NEXT: int _t1 = out[index2];
715+
// CHECK-NEXT: out[index2] += in[index2];
716+
// CHECK-NEXT: int _t2 = out[index3];
717+
// CHECK-NEXT: out[index3] += in[index3];
718+
// CHECK-NEXT: int _t3 = out[index4];
719+
// CHECK-NEXT: out[index4] += in[index4];
720+
// CHECK-NEXT: {
721+
// CHECK-NEXT: out[index4] = _t3;
722+
// CHECK-NEXT: int _r_d3 = _d_out[index4];
723+
// CHECK-NEXT: _d_in[index4] += _r_d3;
724+
// CHECK-NEXT: }
725+
// CHECK-NEXT: {
726+
// CHECK-NEXT: out[index3] = _t2;
727+
// CHECK-NEXT: int _r_d2 = _d_out[index3];
728+
// CHECK-NEXT: _d_in[index3] += _r_d2;
729+
// CHECK-NEXT: }
730+
// CHECK-NEXT: {
731+
// CHECK-NEXT: out[index2] = _t1;
732+
// CHECK-NEXT: int _r_d1 = _d_out[index2];
733+
// CHECK-NEXT: _d_in[index2] += _r_d1;
734+
// CHECK-NEXT: }
735+
// CHECK-NEXT: {
736+
// CHECK-NEXT: out[index1] = _t0;
737+
// CHECK-NEXT: int _r_d0 = _d_out[index1];
738+
// CHECK-NEXT: _d_in[index1] += _r_d0;
739+
// CHECK-NEXT: }
740+
// CHECK-NEXT: }
741+
742+
__global__ void indices_lin_comb(int *out, int *in) {
743+
int index = threadIdx.x + blockIdx.x * blockDim.x;
744+
745+
out[index] += in[2*index];
746+
out[index] += in[1+index];
747+
out[index] += in[threadIdx.x + blockIdx.x * blockDim.x];
748+
out[index] += in[2*(threadIdx.x + blockIdx.x * blockDim.x) + 1];
749+
out[index] += in[1+1];
750+
out[index] += in[index/2];
751+
752+
}
753+
754+
// CHECK: void indices_lin_comb_grad(int *out, int *in, int *_d_out, int *_d_in) {
755+
// CHECK-NEXT: int _d_index = 0;
756+
// CHECK-NEXT: int index0 = threadIdx.x + blockIdx.x * blockDim.x;
757+
// CHECK-NEXT: int _t0 = out[index0];
758+
// CHECK-NEXT: out[index0] += in[2 * index0];
759+
// CHECK-NEXT: int _t1 = out[index0];
760+
// CHECK-NEXT: out[index0] += in[1 + index0];
761+
// CHECK-NEXT: int _t2 = out[index0];
762+
// CHECK-NEXT: out[index0] += in[threadIdx.x + blockIdx.x * blockDim.x];
763+
// CHECK-NEXT: int _t3 = out[index0];
764+
// CHECK-NEXT: unsigned int _t4 = (threadIdx.x + blockIdx.x * blockDim.x);
765+
// CHECK-NEXT: out[index0] += in[2 * _t4 + 1];
766+
// CHECK-NEXT: int _t5 = out[index0];
767+
// CHECK-NEXT: out[index0] += in[1 + 1];
768+
// CHECK-NEXT: int _t6 = out[index0];
769+
// CHECK-NEXT: out[index0] += in[index0 / 2];
770+
// CHECK-NEXT: {
771+
// CHECK-NEXT: out[index0] = _t6;
772+
// CHECK-NEXT: int _r_d5 = _d_out[index0];
773+
// CHECK-NEXT: atomicAdd(&_d_in[index0 / 2], _r_d5);
774+
// CHECK-NEXT: }
775+
// CHECK-NEXT: {
776+
// CHECK-NEXT: out[index0] = _t5;
777+
// CHECK-NEXT: int _r_d4 = _d_out[index0];
778+
// CHECK-NEXT: atomicAdd(&_d_in[1 + 1], _r_d4);
779+
// CHECK-NEXT: }
780+
// CHECK-NEXT: {
781+
// CHECK-NEXT: out[index0] = _t3;
782+
// CHECK-NEXT: int _r_d3 = _d_out[index0];
783+
// CHECK-NEXT: _d_in[2 * _t4 + 1] += _r_d3;
784+
// CHECK-NEXT: }
785+
// CHECK-NEXT: {
786+
// CHECK-NEXT: out[index0] = _t2;
787+
// CHECK-NEXT: int _r_d2 = _d_out[index0];
788+
// CHECK-NEXT: _d_in[threadIdx.x + blockIdx.x * blockDim.x] += _r_d2;
789+
// CHECK-NEXT: }
790+
// CHECK-NEXT: {
791+
// CHECK-NEXT: out[index0] = _t1;
792+
// CHECK-NEXT: int _r_d1 = _d_out[index0];
793+
// CHECK-NEXT: _d_in[1 + index0] += _r_d1;
794+
// CHECK-NEXT: }
795+
// CHECK-NEXT: {
796+
// CHECK-NEXT: out[index0] = _t0;
797+
// CHECK-NEXT: int _r_d0 = _d_out[index0];
798+
// CHECK-NEXT: _d_in[2 * index0] += _r_d0;
799+
// CHECK-NEXT: }
800+
// CHECK-NEXT: }
801+
692802
#define TEST(F, grid, block, shared_mem, use_stream, x, dx, N) \
693803
{ \
694804
int *fives = (int*)malloc(N * sizeof(int)); \
@@ -963,6 +1073,9 @@ int main(void) {
9631073
launch_kernel_4_test.execute(zeros_int, fives_int, 10, out_res, in_res);
9641074
printf("%d, %d, %d\n", in_res[0], in_res[1], in_res[2]); // CHECK-EXEC: 5, 5, 5
9651075

1076+
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
1077+
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
1078+
9661079
free(res);
9671080
free(fives);
9681081
free(zeros);

0 commit comments

Comments
 (0)