Skip to content

Commit dd6f1f4

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 74405b9 commit dd6f1f4

File tree

3 files changed

+254
-24
lines changed

3 files changed

+254
-24
lines changed

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

Lines changed: 126 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,22 +29,42 @@
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+
return I->isMemDepCandidate() ||
64+
(isa<AllocaInst>(I) && cast<AllocaInst>(I)->isUsedWithInAlloca()) ||
65+
I->isStackSaveOrRestoreIntrinsic();
66+
}
67+
4868
Instruction *getInstruction() const { return I; }
4969
void addMemPred(DGNode *PredN) { MemPreds.insert(PredN); }
5070
/// \Returns all memory dependency predecessors.
@@ -53,11 +73,9 @@ class DGNode {
5373
}
5474
/// \Returns true if there is a memory dependency N->this.
5575
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; }
76+
5977
#ifndef NDEBUG
60-
void print(raw_ostream &OS, bool PrintDeps = true) const;
78+
virtual void print(raw_ostream &OS, bool PrintDeps = true) const;
6179
friend raw_ostream &operator<<(DGNode &N, raw_ostream &OS) {
6280
N.print(OS);
6381
return OS;
@@ -66,9 +84,94 @@ class DGNode {
6684
#endif // NDEBUG
6785
};
6886

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

73176
public:
74177
DependencyGraph() {}
@@ -77,10 +180,20 @@ class DependencyGraph {
77180
auto It = InstrToNodeMap.find(I);
78181
return It != InstrToNodeMap.end() ? It->second.get() : nullptr;
79182
}
183+
/// Like getNode() but returns nullptr if \p I is nullptr.
184+
DGNode *getNodeOrNull(Instruction *I) const {
185+
if (I == nullptr)
186+
return nullptr;
187+
return getNode(I);
188+
}
80189
DGNode *getOrCreateNode(Instruction *I) {
81190
auto [It, NotInMap] = InstrToNodeMap.try_emplace(I);
82-
if (NotInMap)
83-
It->second = std::make_unique<DGNode>(I);
191+
if (NotInMap) {
192+
if (I->isMemDepCandidate() || I->isStackSaveOrRestoreIntrinsic())
193+
It->second = std::make_unique<MemDGNode>(I);
194+
else
195+
It->second = std::make_unique<DGNode>(I);
196+
}
84197
return It->second.get();
85198
}
86199
/// 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 InstrInterval &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
InstrInterval DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
3561
if (Instrs.empty())
3662
return {};
@@ -39,10 +65,18 @@ InstrInterval 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
@@ -27,7 +27,7 @@ struct DependencyGraphTest : public testing::Test {
2727
}
2828
};
2929

30-
TEST_F(DependencyGraphTest, DGNode_IsMem) {
30+
TEST_F(DependencyGraphTest, MemDGNode) {
3131
parseIR(C, R"IR(
3232
declare void @llvm.sideeffect()
3333
declare void @llvm.pseudoprobe(i64, i64, i32, i64)
@@ -64,16 +64,16 @@ define void @foo(i8 %v1, ptr %ptr) {
6464

6565
sandboxir::DependencyGraph DAG;
6666
DAG.extend({&*BB->begin(), BB->getTerminator()});
67-
EXPECT_TRUE(DAG.getNode(Store)->isMem());
68-
EXPECT_TRUE(DAG.getNode(Load)->isMem());
69-
EXPECT_FALSE(DAG.getNode(Add)->isMem());
70-
EXPECT_TRUE(DAG.getNode(StackSave)->isMem());
71-
EXPECT_TRUE(DAG.getNode(StackRestore)->isMem());
72-
EXPECT_FALSE(DAG.getNode(SideEffect)->isMem());
73-
EXPECT_FALSE(DAG.getNode(PseudoProbe)->isMem());
74-
EXPECT_TRUE(DAG.getNode(FakeUse)->isMem());
75-
EXPECT_TRUE(DAG.getNode(Call)->isMem());
76-
EXPECT_FALSE(DAG.getNode(Ret)->isMem());
67+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Store)));
68+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Load)));
69+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Add)));
70+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(StackSave)));
71+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(StackRestore)));
72+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(SideEffect)));
73+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(PseudoProbe)));
74+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(FakeUse)));
75+
EXPECT_TRUE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Call)));
76+
EXPECT_FALSE(isa<llvm::sandboxir::MemDGNode>(DAG.getNode(Ret)));
7777
}
7878

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

0 commit comments

Comments
 (0)