Skip to content

[SandboxVec][DAG] Implement PredIterator #111604

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 10, 2024
Merged

[SandboxVec][DAG] Implement PredIterator #111604

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 73 additions & 0 deletions llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,54 @@ enum class DGNodeID {
MemDGNode,
};

class DGNode;
class MemDGNode;
class DependencyGraph;

/// While OpIt points to a Value that is not an Instruction keep incrementing
/// it. \Returns the first iterator that points to an Instruction, or end.
[[nodiscard]] static User::op_iterator skipNonInstr(User::op_iterator OpIt,
User::op_iterator OpItE) {
while (OpIt != OpItE && !isa<Instruction>((*OpIt).get()))
++OpIt;
return OpIt;
}

/// Iterate over both def-use and mem dependencies.
class PredIterator {
User::op_iterator OpIt;
User::op_iterator OpItE;
DenseSet<MemDGNode *>::iterator MemIt;
DGNode *N = nullptr;
DependencyGraph *DAG = nullptr;

PredIterator(const User::op_iterator &OpIt, const User::op_iterator &OpItE,
const DenseSet<MemDGNode *>::iterator &MemIt, DGNode *N,
DependencyGraph &DAG)
: OpIt(OpIt), OpItE(OpItE), MemIt(MemIt), N(N), DAG(&DAG) {}
PredIterator(const User::op_iterator &OpIt, const User::op_iterator &OpItE,
DGNode *N, DependencyGraph &DAG)
: OpIt(OpIt), OpItE(OpItE), N(N), DAG(&DAG) {}
friend class DGNode; // For constructor
friend class MemDGNode; // For constructor

public:
using difference_type = std::ptrdiff_t;
using value_type = DGNode *;
using pointer = value_type *;
using reference = value_type &;
using iterator_category = std::input_iterator_tag;
value_type operator*();
PredIterator &operator++();
PredIterator operator++(int) {
auto Copy = *this;
++(*this);
return Copy;
}
bool operator==(const PredIterator &Other) const;
bool operator!=(const PredIterator &Other) const { return !(*this == Other); }
};

/// A DependencyGraph Node that points to an Instruction and contains memory
/// dependency edges.
class DGNode {
Expand All @@ -63,6 +111,23 @@ class DGNode {
virtual ~DGNode() = default;
/// \Returns true if this is before \p Other in program order.
bool comesBefore(const DGNode *Other) { return I->comesBefore(Other->I); }
using iterator = PredIterator;
virtual iterator preds_begin(DependencyGraph &DAG) {
return PredIterator(skipNonInstr(I->op_begin(), I->op_end()), I->op_end(),
this, DAG);
}
virtual iterator preds_end(DependencyGraph &DAG) {
return PredIterator(I->op_end(), I->op_end(), this, DAG);
}
iterator preds_begin(DependencyGraph &DAG) const {
return const_cast<DGNode *>(this)->preds_begin(DAG);
}
iterator preds_end(DependencyGraph &DAG) const {
return const_cast<DGNode *>(this)->preds_end(DAG);
}
iterator_range<iterator> preds(DependencyGraph &DAG) const {
return make_range(preds_begin(DAG), preds_end(DAG));
}

static bool isStackSaveOrRestoreIntrinsic(Instruction *I) {
if (auto *II = dyn_cast<IntrinsicInst>(I)) {
Expand Down Expand Up @@ -145,6 +210,14 @@ class MemDGNode final : public DGNode {
static bool classof(const DGNode *Other) {
return Other->SubclassID == DGNodeID::MemDGNode;
}
iterator preds_begin(DependencyGraph &DAG) override {
auto OpEndIt = I->op_end();
return PredIterator(skipNonInstr(I->op_begin(), OpEndIt), OpEndIt,
MemPreds.begin(), this, DAG);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it time to move MemPreds into MemDGNode?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, it's in a follow-up patch.

}
iterator preds_end(DependencyGraph &DAG) override {
return PredIterator(I->op_end(), I->op_end(), MemPreds.end(), this, DAG);
}
/// \Returns the previous Mem DGNode in instruction order.
MemDGNode *getPrevNode() const { return PrevMemN; }
/// \Returns the next Mem DGNode in instruction order.
Expand Down
44 changes: 44 additions & 0 deletions llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,54 @@

#include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/SandboxIR/Instruction.h"
#include "llvm/SandboxIR/Utils.h"

namespace llvm::sandboxir {

PredIterator::value_type PredIterator::operator*() {
// If it's a DGNode then we dereference the operand iterator.
if (!isa<MemDGNode>(N)) {
assert(OpIt != OpItE && "Can't dereference end iterator!");
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This confuses me, when N is not a MemDGNode, there is no guarantee OpIt != OpItE.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it's a DGNode then the OpIt != OpItE check doesn't even get executed, it steps right inside the if body.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I refactored this code too to make it more readable.

return DAG->getNode(cast<Instruction>((Value *)*OpIt));
}
// It's a MemDGNode, so we check if we return either the use-def operand,
// or a mem predecessor.
if (OpIt != OpItE)
return DAG->getNode(cast<Instruction>((Value *)*OpIt));
assert(MemIt != cast<MemDGNode>(N)->memPreds().end() &&
"Cant' dereference end iterator!");
return *MemIt;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you simplify this code a bit for just more readability:

 if (!isa<MemDGNode>(N)) {
    assert (OpIt != OpItE && .....);
    return DAG->getNode(cast<Instruction>(*OpIt));
 } 
 if (OpIt != OpItE)
  return DAG->getNode(...);
 // It's a ....

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure I understand how you are going to be using this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The DAG only implements memory dependencies, but not use-def dependencies to save memory. This iterator hides this detail and provides a way to access all dependencies regardless of whether they are use-def or mem. This is useful for the scheduler who doesn't care about the nature of the dependency.


PredIterator &PredIterator::operator++() {
// If it's a DGNode then we increment the use-def iterator.
if (!isa<MemDGNode>(N)) {
assert(OpIt != OpItE && "Already at end!");
++OpIt;
// Skip operands that are not instructions.
OpIt = skipNonInstr(OpIt, OpItE);
return *this;
}
// It's a MemDGNode, so if we are not at the end of the use-def iterator we
// need to first increment that.
if (OpIt != OpItE) {
++OpIt;
// Skip operands that are not instructions.
OpIt = skipNonInstr(OpIt, OpItE);
return *this;
}
assert(MemIt != cast<MemDGNode>(N)->memPreds().end() && "Already at end!");
++MemIt;
return *this;
}

bool PredIterator::operator==(const PredIterator &Other) const {
assert(DAG == Other.DAG && "Iterators of different DAGs!");
assert(N == Other.N && "Iterators of different nodes!");
return OpIt == Other.OpIt && MemIt == Other.MemIt;
}

#ifndef NDEBUG
void DGNode::print(raw_ostream &OS, bool PrintDeps) const {
I->dumpOS(OS);
Expand Down
41 changes: 41 additions & 0 deletions llvm/unittests/Transforms/Vectorize/SandboxVectorizer/DependencyGraphTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,12 +240,53 @@ define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
EXPECT_TRUE(N1->hasMemPred(N0));
EXPECT_FALSE(N0->hasMemPred(N1));

// Check preds().
EXPECT_TRUE(N0->preds(DAG).empty());
EXPECT_THAT(N1->preds(DAG), testing::ElementsAre(N0));

// Check memPreds().
EXPECT_TRUE(N0->memPreds().empty());
EXPECT_THAT(N1->memPreds(), testing::ElementsAre(N0));
EXPECT_TRUE(N2->memPreds().empty());
}

TEST_F(DependencyGraphTest, Preds) {
parseIR(C, R"IR(
declare ptr @bar(i8)
define i8 @foo(i8 %v0, i8 %v1) {
%add0 = add i8 %v0, %v0
%add1 = add i8 %v1, %v1
%add2 = add i8 %add0, %add1
%ptr = call ptr @bar(i8 %add1)
store i8 %add2, ptr %ptr
ret i8 %add2
}
)IR");
llvm::Function *LLVMF = &*M->getFunction("foo");
sandboxir::Context Ctx(C);
auto *F = Ctx.createFunction(LLVMF);
auto *BB = &*F->begin();
auto It = BB->begin();
sandboxir::DependencyGraph DAG(getAA(*LLVMF));
DAG.extend({&*BB->begin(), BB->getTerminator()});

auto *AddN0 = DAG.getNode(cast<sandboxir::BinaryOperator>(&*It++));
auto *AddN1 = DAG.getNode(cast<sandboxir::BinaryOperator>(&*It++));
auto *AddN2 = DAG.getNode(cast<sandboxir::BinaryOperator>(&*It++));
auto *CallN = DAG.getNode(cast<sandboxir::CallInst>(&*It++));
auto *StN = DAG.getNode(cast<sandboxir::StoreInst>(&*It++));
auto *RetN = DAG.getNode(cast<sandboxir::ReturnInst>(&*It++));

// Check preds().
EXPECT_THAT(AddN0->preds(DAG), testing::ElementsAre());
EXPECT_THAT(AddN1->preds(DAG), testing::ElementsAre());
EXPECT_THAT(AddN2->preds(DAG), testing::ElementsAre(AddN0, AddN1));
EXPECT_THAT(CallN->preds(DAG), testing::ElementsAre(AddN1));
EXPECT_THAT(StN->preds(DAG),
testing::UnorderedElementsAre(CallN, CallN, AddN2));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

adding a note somewhere that this can return the same instruction multiple times would be good

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah there is no filtering out of duplicates. I will add a note in the description.

EXPECT_THAT(RetN->preds(DAG), testing::ElementsAre(AddN2));
}

TEST_F(DependencyGraphTest, MemDGNode_getPrevNode_getNextNode) {
parseIR(C, R"IR(
define void @foo(ptr %ptr, i8 %v0, i8 %v1) {
Expand Down
Loading