Visitor: Add VisitorStrategy::ReversePostOrder

This will be convenient for certain lowering passes that depend on
substituting values in a more complex pattern than just via
replaceAllUsesWith.

Author: nhaehnle

example/ExampleMain.cpp

@@ -58,6 +58,9 @@ cl::list<std::string> g_inputs(cl::Positional, cl::ZeroOrMore,
 cl::opt<bool> g_typedPointers("typed-pointers", cl::init(false),
                                cl::desc("Disable opaque pointers"));
 
+cl::opt<bool> g_rpot("rpot", cl::init(false),
+                     cl::desc("Visit functions in reverse post-order"));
+
 /// Test intrinsics that are generic over unnamed struct types.
 void useUnnamedStructTypes(Builder &b) {
   StructType *t1 = StructType::create({b.getInt32Ty()}, "");
@@ -156,8 +159,15 @@ struct llvm_dialects::VisitorPayloadProjection<VisitorNest, raw_ostream> {
 LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD(VisitorContainer, nest)
 LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD(VisitorNest, inner)
 
-void exampleVisit(Module &module) {
-  auto visitor =
+// Get a visitor with RPOT or default strategy.
+//
+// Using a template in this way is a bit convoluted, but the point here is to
+// demonstrate the pattern of constructing a visitor as a function-local
+// static variable whose initialization doesn't refer to any runtime values,
+// i.e. if C++ had a strong enough compile-time evaluation (constexpr), it
+// should be possible to evaluate the initialization entirely at compile-time.
+template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
+  static const auto visitor =
       VisitorBuilder<VisitorContainer>()
           .nest<VisitorNest>([](VisitorBuilder<VisitorNest> &b) {
             b.add<xd::ReadOp>([](VisitorNest &self, xd::ReadOp &op) {
@@ -173,7 +183,15 @@ void exampleVisit(Module &module) {
                                      xd::ITruncOp &op) { inner.counter++; });
             });
           })
+          .setStrategy(rpot ? VisitorStrategy::ReversePostOrder
+                            : VisitorStrategy::Default)
           .build();
+  return visitor;
+}
+
+void exampleVisit(Module &module) {
+  const auto &visitor =
+      g_rpot ? getExampleVisitor<true>() : getExampleVisitor<false>();
 
   VisitorContainer container;
   container.nest.out = &outs();

include/llvm-dialects/Dialect/Visitor.h

@@ -32,11 +32,17 @@ namespace llvm_dialects {
 template <typename PayloadT>
 class Visitor;
 
-/// The strategy for implementation a visitor of dialect ops, 
+/// The iteration strategy of Visitor.
 enum class VisitorStrategy {
+  /// Pick a reasonable default.
+  Default,
+
   /// Iterate over all instructions to find the ones that match.
   ByInstruction,
 
+  /// Iterate over all instructions in reverse post-order.
+  ReversePostOrder,
+
   /// Iterate over the function declarations in a module, filter out those for
   /// relevant dialect ops, and then iterate over their users.
   ///
@@ -116,7 +122,7 @@ class VisitorBuilderBase {
   explicit VisitorBuilderBase(VisitorBuilderBase *parent) : m_parent(parent) {}
   ~VisitorBuilderBase();
 
-  void setStrategy(VisitorStrategy strategy) { m_strategy = strategy; }
+  void setStrategy(VisitorStrategy strategy);
 
   void add(const OpDescription &desc, void *extra, VisitorCallback *fn);
 
@@ -126,7 +132,7 @@ class VisitorBuilderBase {
 
 private:
   VisitorBuilderBase *m_parent = nullptr;
-  VisitorStrategy m_strategy = VisitorStrategy::ByFunctionDeclaration;
+  VisitorStrategy m_strategy = VisitorStrategy::Default;
   llvm::SmallVector<VisitorCase> m_cases;
   llvm::SmallVector<PayloadProjection> m_projections;
 };
@@ -155,7 +161,8 @@ class VisitorBase {
 /// See @ref VisitorBuilder for how to instantiate this class.
 ///
 /// Instructions are not necessarily visited in the order in which they appear
-/// in the containing function or even basic block.
+/// in the containing function or even basic block. If a fixed order is
+/// required, select the ReversePostOrder strategy at VisitorBuilder time.
 ///
 /// Callbacks must not delete or remove their instruction argument.
 template <typename PayloadT>

lib/Dialect/Visitor.cpp

@@ -16,6 +16,7 @@
 
 #include "llvm-dialects/Dialect/Visitor.h"
 
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
@@ -60,9 +61,20 @@ VisitorBuilderBase::~VisitorBuilderBase() {
 
     // Copy all cases to the root.
     ancestor->m_cases.append(m_cases.begin(), m_cases.end());
+
+    ancestor->setStrategy(m_strategy);
   }
 }
 
+void VisitorBuilderBase::setStrategy(VisitorStrategy strategy) {
+  if (strategy == VisitorStrategy::Default)
+    return;
+
+  assert(m_strategy == VisitorStrategy::Default || m_strategy == strategy);
+
+  m_strategy = strategy;
+}
+
 void VisitorBuilderBase::add(const OpDescription &desc, void *extra, VisitorCallback *fn) {
   VisitorCase theCase;
   theCase.description = &desc;
@@ -76,6 +88,9 @@ VisitorBase::VisitorBase(VisitorBuilderBase &&builder)
     : m_strategy(builder.m_strategy), m_cases(std::move(builder.m_cases)),
       m_projections(std::move(builder.m_projections)) {
   assert(!builder.m_parent);
+
+  if (m_strategy == VisitorStrategy::Default)
+    m_strategy = VisitorStrategy::ByFunctionDeclaration;
 }
 
 void VisitorBase::call(const VisitorCase &theCase, void *payload,
@@ -109,6 +124,15 @@ void VisitorBase::visit(void *payload, Function &fn) const {
     return;
   }
 
+  if (m_strategy == VisitorStrategy::ReversePostOrder) {
+    ReversePostOrderTraversal<Function *> rpot(&fn);
+    for (BasicBlock *bb : rpot) {
+      for (Instruction &inst : *bb)
+        visit(payload, inst);
+    }
+    return;
+  }
+
   for (Function &decl : fn.getParent()->functions()) {
     if (!decl.isDeclaration())
       continue;
@@ -133,7 +157,7 @@ void VisitorBase::visit(void *payload, Function &fn) const {
 }
 
 void VisitorBase::visit(void *payload, Module &module) const {
-  if (m_strategy == VisitorStrategy::ByInstruction) {
+  if (m_strategy != VisitorStrategy::ByFunctionDeclaration) {
     for (Function &fn : module.functions()) {
       if (!fn.isDeclaration())
         visit(payload, fn);

test/example/visitor-rpot.ll

@@ -0,0 +1,26 @@
+; RUN: llvm-dialects-example -visit %s | FileCheck --check-prefixes=DEFAULT %s
+; RUN: llvm-dialects-example -visit -rpot %s | FileCheck --check-prefixes=RPOT %s
+
+; DEFAULT: visiting WriteOp: call void (...) @xd.write(i32 0)
+; DEFAULT: visiting WriteOp: call void (...) @xd.write(i32 2)
+; DEFAULT: visiting WriteOp: call void (...) @xd.write(i32 1)
+
+; RPOT: visiting WriteOp: call void (...) @xd.write(i32 0)
+; RPOT: visiting WriteOp: call void (...) @xd.write(i32 1)
+; RPOT: visiting WriteOp: call void (...) @xd.write(i32 2)
+
+define void @test1(ptr %p) {
+entry:
+  call void (...) @xd.write(i32 0)
+  br label %a
+
+b:
+  call void (...) @xd.write(i32 2)
+  ret void
+
+a:
+  call void (...) @xd.write(i32 1)
+  br label %b
+}
+
+declare void @xd.write(...)