unroll: Patch phi nodes inside loops
MODIFY
ir/function.cpp
ir/function.cpp
@@ -58,6 +58,45 @@ void BasicBlock::replaceTargetWith(const BasicBlock *from,
}
}
+void BasicBlock::replacePhiSrcWith(const string &from,
+ const string &to) {
+ for (auto &i : instrs()) {
+ if (Phi *phi = dynamic_cast<Phi*>(const_cast<Instr*>(&i)))
+ phi->replaceSourceWith(from, to);
+ }
+}
+
+void BasicBlock::removePhiSrc(const BasicBlock *bb) {
+ for (auto &i : instrs()) {
+ if (Phi *phi = dynamic_cast<Phi*>(const_cast<Instr*>(&i))) {
+ phi->removeValue(bb->getName());
+ }
+ }
+}
+
+void BasicBlock::removePhiSrcsExcept(const BasicBlock *bb) {
+ for (auto &i : instrs()) {
+ if (Phi *phi = dynamic_cast<Phi*>(const_cast<Instr*>(&i))) {
+ for (auto &src : phi->sources())
+ if (bb->getName() != src)
+ phi->removeValue(src);
+ }
+ }
+}
+
+vector<string> BasicBlock::phiSources() const {
+ vector<string> s;
+
+ for (auto &i : instrs()) {
+ if (Phi *phi = dynamic_cast<Phi*>(const_cast<Instr*>(&i))) {
+ auto v = phi->sources();
+ s.insert(s.end(), v.begin(), v.end());
+ }
+ }
+
+ return s;
+}
+
unique_ptr<BasicBlock> BasicBlock::dup(const string &suffix) const {
auto newbb = make_unique<BasicBlock>(name + suffix);
for (auto &i : instrs()) {
@@ -408,6 +447,19 @@ void Function::unroll(unsigned k) {
bbmap[bb].emplace_back(bb);
}
+ // find latch block
+ // TODO: handle the case where multiple loops share one header
+ const BasicBlock *latch = nullptr;
+ for (auto &[bb, _] : bbmap) {
+ for (auto &tgt : bb->targets()) {
+ if (header == &tgt) {
+ latch = bb;
+ break;
+ }
+ }
+ if (latch) break;
+ }
+
// Clone BBs
// Note that the BBs list must be iterated in top-sort order so that
// values from previous BBs are available in vmap
@@ -462,7 +514,33 @@ void Function::unroll(unsigned k) {
}
}
- // TODO: handle phis
+ // Patch phi sources
+ // for header, remove the phi source that comes from backedge
+ header->removePhiSrc(latch);
+ // for the copies of header
+ auto header_copies = bbmap[header];
+ for (unsigned unroll = 1, e = header_copies.size(); unroll < e; ++unroll) {
+ auto *cloned_header = header_copies[unroll];
+ // preserve only latch as source
+ cloned_header->removePhiSrcsExcept(latch);
+ // rewrite incoming phi source in Kth copy's header to (K-1)th copy's latch
+ auto to = bbmap[latch][unroll-1]->getName();
+ cloned_header->replacePhiSrcWith(latch->getName(), to);
+ }
+
+ // for non-headers, update phi sources that is coming from inside loop
+ for (auto &[bb, copies] : bbmap) {
+ if (bb == header)
+ continue;
+ for (unsigned unroll = 1, e = copies.size(); unroll < e; ++unroll) {
+ auto *cloned = copies[unroll];
+ auto srcs = bb->phiSources();
+ for (auto from : srcs) {
+ auto to = bbmap[&getBB(from)][unroll]->getName();
+ cloned->replacePhiSrcWith(from, to);
+ }
+ }
+ }
}
topSort();

MODIFY
ir/function.h
ir/function.h
@@ -47,6 +47,10 @@ public:
bool empty() const { return m_instrs.empty(); }
JumpInstr::it_helper targets() const;
void replaceTargetWith(const BasicBlock *from, const BasicBlock *to);
+ void replacePhiSrcWith(const std::string &from, const std::string &to);
+ void removePhiSrc(const BasicBlock *bb);
+ void removePhiSrcsExcept(const BasicBlock *bb);
+ std::vector<std::string> phiSources() const;
std::unique_ptr<BasicBlock> dup(const std::string &suffix) const;

MODIFY
ir/instr.cpp
ir/instr.cpp
@@ -2047,6 +2047,23 @@ void Phi::removeValue(const string &BB_name) {
}
}
+vector<string> Phi::sources() const {
+ vector<string> s;
+ for (auto &[_, bb] : values) {
+ s.emplace_back(bb);
+ }
+ return s;
+}
+
+void Phi::replaceSourceWith(const string &from, const string &to) {
+ for (auto &[_, bb] : values) {
+ if (bb == from) {
+ bb = to;
+ break;
+ }
+ }
+}
+
vector<Value*> Phi::operands() const {
vector<Value*> v;
for (auto &[val, bb] : values) {

MODIFY
ir/instr.h
ir/instr.h
@@ -294,6 +294,8 @@ public:
void addValue(Value &val, std::string &&BB_name);
void removeValue(const std::string &BB_name);
+ void replaceSourceWith(const std::string &from, const std::string &to);
+ std::vector<std::string> sources() const;
std::vector<Value*> operands() const override;
void rauw(const Value &what, Value &with) override;

unroll: Patch phi nodes inside loops (#552)