Skip to content

Commit 098cc79

Browse files
committed
[SandboxVec][DAG] MemDGNode for memory-dependency candidate nodes
This patch implements the MemDGNode class for DAG nodes that are candidates for memory dependencies. These nodes form a chain that is accessible by `getPrevNode()` and `getNextNode()`. The patch also implements the MemDGNodeIterator for iterating over the MemDGNode chain. Finally, makeMemRange() returns a MemDGNodeIterator range given an instruction interval of memory or non-memory instructions. This will be used in a follow-up patch for dependency scanning.
1 parent 1fb3fe0 commit 098cc79

File tree

3 files changed

+256
-24
lines changed

3 files changed

+256
-24
lines changed

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h

Lines changed: 128 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,22 +29,44 @@
2929

3030
namespace llvm::sandboxir {
3131

32+
class DependencyGraph;
33+
class MemDGNode;
34+
35+
/// SubclassIDs for isa/dyn_cast etc.
36+
enum class DGNodeID {
37+
DGNode,
38+
MemDGNode,
39+
};
40+
3241
/// A DependencyGraph Node that points to an Instruction and contains memory
3342
/// dependency edges.
3443
class DGNode {
44+
protected:
3545
Instruction *I;
46+
// TODO: Use a PointerIntPair for SubclassID and I.
47+
/// For isa/dyn_cast etc.
48+
DGNodeID SubclassID;
3649
/// Memory predecessors.
3750
DenseSet<DGNode *> MemPreds;
38-
/// This is true if this may read/write memory, or if it has some ordering
39-
/// constraints, like with stacksave/stackrestore and alloca/inalloca.
40-
bool IsMem;
51+
52+
DGNode(Instruction *I, DGNodeID ID) : I(I), SubclassID(ID) {}
53+
friend class MemDGNode; // For constructor.
4154

4255
public:
43-
DGNode(Instruction *I) : I(I) {
44-
IsMem = I->isMemDepCandidate() ||
45-
(isa<AllocaInst>(I) && cast<AllocaInst>(I)->isUsedWithInAlloca()) ||
46-
I->isStackSaveOrRestoreIntrinsic();
56+
DGNode(Instruction *I) : I(I), SubclassID(DGNodeID::DGNode) {
57+
assert(!isMemDepCandidate(I) && "Expected Non-Mem instruction, ");
4758
}
59+
DGNode(const DGNode &Other) = delete;
60+
virtual ~DGNode() = default;
61+
/// \Returns true if \p I is a memory dependency candidate instruction.
62+
static bool isMemDepCandidate(Instruction *I) {
63+
AllocaInst *Alloca;
64+
return I->isMemDepCandidate() ||
65+
((Alloca = dyn_cast<AllocaInst>(I)) &&
66+
Alloca->isUsedWithInAlloca()) ||
67+
I->isStackSaveOrRestoreIntrinsic();
68+
}
69+
4870
Instruction *getInstruction() const { return I; }
4971
void addMemPred(DGNode *PredN) { MemPreds.insert(PredN); }
5072
/// \Returns all memory dependency predecessors.
@@ -53,11 +75,9 @@ class DGNode {
5375
}
5476
/// \Returns true if there is a memory dependency N->this.
5577
bool hasMemPred(DGNode *N) const { return MemPreds.count(N); }
56-
/// \Returns true if this may read/write memory, or if it has some ordering
57-
/// constraints, like with stacksave/stackrestore and alloca/inalloca.
58-
bool isMem() const { return IsMem; }
78+
5979
#ifndef NDEBUG
60-
void print(raw_ostream &OS, bool PrintDeps = true) const;
80+
virtual void print(raw_ostream &OS, bool PrintDeps = true) const;
6181
friend raw_ostream &operator<<(DGNode &N, raw_ostream &OS) {
6282
N.print(OS);
6383
return OS;
@@ -66,9 +86,94 @@ class DGNode {
6686
#endif // NDEBUG
6787
};
6888

89+
/// A DependencyGraph Node for instructiosn that may read/write memory, or have
90+
/// some ordering constraints, like with stacksave/stackrestore and
91+
/// alloca/inalloca.
92+
class MemDGNode final : public DGNode {
93+
MemDGNode *PrevMemN = nullptr;
94+
MemDGNode *NextMemN = nullptr;
95+
96+
void setNextNode(MemDGNode *N) { NextMemN = N; }
97+
void setPrevNode(MemDGNode *N) { PrevMemN = N; }
98+
friend class DependencyGraph; // For setNextNode(), setPrevNode().
99+
100+
public:
101+
MemDGNode(Instruction *I) : DGNode(I, DGNodeID::MemDGNode) {
102+
assert(isMemDepCandidate(I) && "Expected Mem instruction!");
103+
}
104+
static bool classof(const DGNode *Other) {
105+
return Other->SubclassID == DGNodeID::MemDGNode;
106+
}
107+
/// \Returns the previous Mem DGNode in instruction order.
108+
MemDGNode *getPrevNode() const { return PrevMemN; }
109+
/// \Returns the next Mem DGNode in instruction order.
110+
MemDGNode *getNextNode() const { return NextMemN; }
111+
};
112+
113+
/// Walks in the order of the instruction chain but skips non-mem Nodes.
114+
/// This is used for building/updating the DAG.
115+
class MemDGNodeIterator {
116+
MemDGNode *N;
117+
118+
public:
119+
using difference_type = std::ptrdiff_t;
120+
using value_type = DGNode;
121+
using pointer = value_type *;
122+
using reference = value_type &;
123+
using iterator_category = std::bidirectional_iterator_tag;
124+
MemDGNodeIterator(MemDGNode *N) : N(N) {}
125+
MemDGNodeIterator &operator++() {
126+
assert(N != nullptr && "Already at end!");
127+
N = N->getNextNode();
128+
return *this;
129+
}
130+
MemDGNodeIterator operator++(int) {
131+
auto ItCopy = *this;
132+
++*this;
133+
return ItCopy;
134+
}
135+
MemDGNodeIterator &operator--() {
136+
N = N->getPrevNode();
137+
return *this;
138+
}
139+
MemDGNodeIterator operator--(int) {
140+
auto ItCopy = *this;
141+
--*this;
142+
return ItCopy;
143+
}
144+
pointer operator*() { return N; }
145+
const DGNode *operator*() const { return N; }
146+
bool operator==(const MemDGNodeIterator &Other) const { return N == Other.N; }
147+
bool operator!=(const MemDGNodeIterator &Other) const {
148+
return !(*this == Other);
149+
}
150+
};
151+
152+
/// A MemDGNodeIterator with convenience builders and dump().
153+
class DGNodeRange : public iterator_range<MemDGNodeIterator> {
154+
public:
155+
DGNodeRange(MemDGNodeIterator Begin, MemDGNodeIterator End)
156+
: iterator_range(Begin, End) {}
157+
/// An empty range.
158+
DGNodeRange()
159+
: iterator_range(MemDGNodeIterator(nullptr), MemDGNodeIterator(nullptr)) {
160+
}
161+
/// Given \p Instrs it finds their closest mem nodes in the interval and
162+
/// returns the corresponding mem range. Note: BotN (or its neighboring mem
163+
/// node) is included in the range.
164+
static DGNodeRange makeMemRange(const Interval<Instruction> &Instrs,
165+
DependencyGraph &DAG);
166+
static DGNodeRange makeEmptyMemRange() { return DGNodeRange(); }
167+
#ifndef NDEBUG
168+
LLVM_DUMP_METHOD void dump() const;
169+
#endif
170+
};
171+
69172
class DependencyGraph {
70173
private:
71174
DenseMap<Instruction *, std::unique_ptr<DGNode>> InstrToNodeMap;
175+
/// The DAG spans across all instructions in this interval.
176+
Interval<Instruction> DAGInterval;
72177

73178
public:
74179
DependencyGraph() {}
@@ -77,10 +182,20 @@ class DependencyGraph {
77182
auto It = InstrToNodeMap.find(I);
78183
return It != InstrToNodeMap.end() ? It->second.get() : nullptr;
79184
}
185+
/// Like getNode() but returns nullptr if \p I is nullptr.
186+
DGNode *getNodeOrNull(Instruction *I) const {
187+
if (I == nullptr)
188+
return nullptr;
189+
return getNode(I);
190+
}
80191
DGNode *getOrCreateNode(Instruction *I) {
81192
auto [It, NotInMap] = InstrToNodeMap.try_emplace(I);
82-
if (NotInMap)
83-
It->second = std::make_unique<DGNode>(I);
193+
if (NotInMap) {
194+
if (I->isMemDepCandidate() || I->isStackSaveOrRestoreIntrinsic())
195+
It->second = std::make_unique<MemDGNode>(I);
196+
else
197+
It->second = std::make_unique<DGNode>(I);
198+
}
84199
return It->second.get();
85200
}
86201
/// Build/extend the dependency graph such that it includes \p Instrs. Returns

llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,32 @@ void DGNode::dump() const {
3131
}
3232
#endif // NDEBUG
3333

34+
DGNodeRange DGNodeRange::makeMemRange(const Interval<Instruction> &Instrs,
35+
DependencyGraph &DAG) {
36+
// If top or bottom instructions are not mem-dep candidate nodes we need to
37+
// walk down/up the chain and find the mem-dep ones.
38+
Instruction *MemTopI = Instrs.top();
39+
Instruction *MemBotI = Instrs.bottom();
40+
while (!DGNode::isMemDepCandidate(MemTopI) && MemTopI != MemBotI)
41+
MemTopI = MemTopI->getNextNode();
42+
while (!DGNode::isMemDepCandidate(MemBotI) && MemBotI != MemTopI)
43+
MemBotI = MemBotI->getPrevNode();
44+
// If we couldn't find a mem node in range TopN - BotN then it's empty.
45+
if (!DGNode::isMemDepCandidate(MemTopI))
46+
return {};
47+
// Now that we have the mem-dep nodes, create and return the range.
48+
return DGNodeRange(
49+
MemDGNodeIterator(cast<MemDGNode>(DAG.getNode(MemTopI))),
50+
MemDGNodeIterator(cast<MemDGNode>(DAG.getNode(MemBotI))->getNextNode()));
51+
}
52+
53+
#ifndef NDEBUG
54+
void DGNodeRange::dump() const {
55+
for (const DGNode *N : *this)
56+
N->dump();
57+
}
58+
#endif // NDEBUG
59+
3460
Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
3561
if (Instrs.empty())
3662
return {};
@@ -39,10 +65,18 @@ Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
3965
auto *TopI = Interval.top();
4066
auto *BotI = Interval.bottom();
4167
DGNode *LastN = getOrCreateNode(TopI);
68+
MemDGNode *LastMemN = dyn_cast<MemDGNode>(LastN);
4269
for (Instruction *I = TopI->getNextNode(), *E = BotI->getNextNode(); I != E;
4370
I = I->getNextNode()) {
4471
auto *N = getOrCreateNode(I);
4572
N->addMemPred(LastN);
73+
// Build the Mem node chain.
74+
if (auto *MemN = dyn_cast<MemDGNode>(N)) {
75+
MemN->setPrevNode(LastMemN);
76+
if (LastMemN != nullptr)
77+
LastMemN->setNextNode(MemN);
78+
LastMemN = MemN;
79+
}
4680
LastN = N;
4781
}
4882
return Interval;

llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp

Lines changed: 94 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ struct DependencyGraphTest : public testing::Test {
2929
}
3030
};
3131

32-
TEST_F(DependencyGraphTest, DGNode_IsMem) {
32+
TEST_F(DependencyGraphTest, MemDGNode) {
3333
parseIR(C, R"IR(
3434
declare void @llvm.sideeffect()
3535
declare void @llvm.pseudoprobe(i64, i64, i32, i64)
@@ -66,16 +66,16 @@ define void @foo(i8 %v1, ptr %ptr) {
6666

6767
sandboxir::DependencyGraph DAG;
6868
DAG.extend({&*BB->begin(), BB->getTerminator()});
69-
EXPECT_TRUE(DAG.getNode(Store)->isMem());
70-
EXPECT_TRUE(DAG.getNode(Load)->isMem());
71-
EXPECT_FALSE(DAG.getNode(Add)->isMem());
72-
EXPECT_TRUE(DAG.getNode(StackSave)->isMem());
73-
EXPECT_TRUE(DAG.getNode(StackRestore)->isMem());
74-
EXPECT_FALSE(DAG.getNode(SideEffect)->isMem());
75-
EXPECT_FALSE(DAG.getNode(PseudoProbe)->isMem());
76-
EXPECT_TRUE(DAG.getNode(FakeUse)->isMem());
77-
EXPECT_TRUE(DAG.getNode(Call)->isMem());
78-
EXPECT_FALSE(DAG.getNode(Ret)->isMem());
69+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Store)));
70+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Load)));
71+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Add)));
72+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(StackSave)));
73+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(StackRestore)));
74+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(SideEffect)));
75+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(PseudoProbe)));
76+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(FakeUse)));
77+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Call)));
78+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Ret)));
7979
}
8080

8181
TEST_F(DependencyGraphTest, Basic) {
@@ -115,3 +115,86 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
115115
EXPECT_THAT(N1->memPreds(), testing::ElementsAre(N0));
116116
EXPECT_THAT(N2->memPreds(), testing::ElementsAre(N1));
117117
}
118+
119+
TEST_F(DependencyGraphTest, MemDGNode_getPrevNode_getNextNode) {
120+
parseIR(C, R"IR(
121+
define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
122+
store i8 %v0, ptr %ptr
123+
add i8 %v0, %v0
124+
store i8 %v1, ptr %ptr
125+
ret void
126+
}
127+
)IR");
128+
llvm::Function *LLVMF = &*M->getFunction("foo");
129+
sandboxir::Context Ctx(C);
130+
auto *F = Ctx.createFunction(LLVMF);
131+
auto *BB = &*F->begin();
132+
auto It = BB->begin();
133+
auto *S0 = cast<sandboxir::StoreInst>(&*It++);
134+
[[maybe_unused]] auto *Add = cast<sandboxir::BinaryOperator>(&*It++);
135+
auto *S1 = cast<sandboxir::StoreInst>(&*It++);
136+
[[maybe_unused]] auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
137+
138+
sandboxir::DependencyGraph DAG;
139+
DAG.extend({&*BB->begin(), BB->getTerminator()});
140+
141+
auto *S0N = cast<sandboxir::MemDGNode>(DAG.getNode(S0));
142+
auto *S1N = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
143+
144+
EXPECT_EQ(S0N->getPrevNode(), nullptr);
145+
EXPECT_EQ(S0N->getNextNode(), S1N);
146+
147+
EXPECT_EQ(S1N->getPrevNode(), S0N);
148+
EXPECT_EQ(S1N->getNextNode(), nullptr);
149+
}
150+
151+
TEST_F(DependencyGraphTest, DGNodeRange) {
152+
parseIR(C, R"IR(
153+
define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
154+
add i8 %v0, %v0
155+
store i8 %v0, ptr %ptr
156+
add i8 %v0, %v0
157+
store i8 %v1, ptr %ptr
158+
ret void
159+
}
160+
)IR");
161+
llvm::Function *LLVMF = &*M->getFunction("foo");
162+
sandboxir::Context Ctx(C);
163+
auto *F = Ctx.createFunction(LLVMF);
164+
auto *BB = &*F->begin();
165+
auto It = BB->begin();
166+
auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);
167+
auto *S0 = cast<sandboxir::StoreInst>(&*It++);
168+
auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);
169+
auto *S1 = cast<sandboxir::StoreInst>(&*It++);
170+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
171+
172+
sandboxir::DependencyGraph DAG;
173+
DAG.extend({&*BB->begin(), BB->getTerminator()});
174+
175+
auto *S0N = cast<sandboxir::MemDGNode>(DAG.getNode(S0));
176+
auto *S1N = cast<sandboxir::MemDGNode>(DAG.getNode(S1));
177+
178+
// Check empty range.
179+
EXPECT_THAT(sandboxir::DGNodeRange::makeEmptyMemRange(),
180+
testing::ElementsAre());
181+
182+
// Both TopN and BotN are memory.
183+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({S0, S1}, DAG),
184+
testing::ElementsAre(S0N, S1N));
185+
// Only TopN is memory.
186+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({S0, Ret}, DAG),
187+
testing::ElementsAre(S0N, S1N));
188+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({S0, Add1}, DAG),
189+
testing::ElementsAre(S0N));
190+
// Only BotN is memory.
191+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({Add0, S1}, DAG),
192+
testing::ElementsAre(S0N, S1N));
193+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({Add0, S0}, DAG),
194+
testing::ElementsAre(S0N));
195+
// Neither TopN or BotN is memory.
196+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({Add0, Ret}, DAG),
197+
testing::ElementsAre(S0N, S1N));
198+
EXPECT_THAT(sandboxir::DGNodeRange::makeMemRange({Add0, Add0}, DAG),
199+
testing::ElementsAre());
200+
}

0 commit comments

Comments
 (0)