Refactor the architecture of mlir-reduce

Add iterator for ReductionNode traversal and use range to indicate the region we would like to keep. Refactor the interaction between Pass/Tester/ReductionNode.
Now it'll be easier to add new traversal type and OpReducer

Reviewed By: jpienaar, rriddle

Differential Revision: https://reviews.llvm.org/D99713

Author: jpienaar

mlir/include/mlir/Reducer/Passes/OpReducer.h

@@ -15,65 +15,52 @@
 #ifndef MLIR_REDUCER_PASSES_OPREDUCER_H
 #define MLIR_REDUCER_PASSES_OPREDUCER_H
 
-#include "mlir/IR/Region.h"
+#include <limits>
+
 #include "mlir/Reducer/ReductionNode.h"
-#include "mlir/Reducer/ReductionTreeUtils.h"
 #include "mlir/Reducer/Tester.h"
 
 namespace mlir {
 
-class OpReducerImpl {
+class OpReducer {
 public:
-  OpReducerImpl(
-      llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps);
-
-  /// Return the name of this reducer class.
-  StringRef getName();
-
-  /// Return the initial transformSpace containing the transformable indices.
-  std::vector<bool> initTransformSpace(ModuleOp module);
-
-  /// Generate variants by removing OpType operations from the module in the
-  /// parent and link the variants as childs in the Reduction Tree Pass.
-  void generateVariants(ReductionNode *parent, const Tester &test,
-                        int numVariants);
-
-  /// Generate variants by removing OpType operations from the module in the
-  /// parent and link the variants as childs in the Reduction Tree Pass. The
-  /// transform argument defines the function used to remove the OpTpye
-  /// operations in range of indexed OpType operations.
-  void generateVariants(ReductionNode *parent, const Tester &test,
-                        int numVariants,
-                        llvm::function_ref<void(ModuleOp, int, int)> transform);
-
-private:
-  llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps;
+  virtual ~OpReducer() = default;
+  /// According to rangeToKeep, try to reduce the given module. We implicitly
+  /// number each interesting operation and rangeToKeep indicates that if an
+  /// operation's number falls into certain range, then we will not try to
+  /// reduce that operation.
+  virtual void reduce(ModuleOp module,
+                      ArrayRef<ReductionNode::Range> rangeToKeep) = 0;
+  /// Return the number of certain kind of operations that we would like to
+  /// reduce. This can be used to build a range map to exclude uninterested
+  /// operations.
+  virtual int getNumTargetOps(ModuleOp module) const = 0;
 };
 
-/// The OpReducer class defines a variant generator method that produces
-/// multiple variants by eliminating different OpType operations from the
-/// parent module.
+/// Reducer is a helper class to remove potential uninteresting operations from
+/// module.
 template <typename OpType>
-class OpReducer {
+class Reducer : public OpReducer {
 public:
-  OpReducer() : impl(new OpReducerImpl(getSpecificOps)) {}
+  ~Reducer() override = default;
 
-  /// Returns the vector of pointer to the OpType operations in the module.
-  static std::vector<Operation *> getSpecificOps(ModuleOp module) {
-    std::vector<Operation *> ops;
-    for (auto op : module.getOps<OpType>()) {
-      ops.push_back(op);
-    }
-    return ops;
+  int getNumTargetOps(ModuleOp module) const override {
+    return std::distance(module.getOps<OpType>().begin(),
+                         module.getOps<OpType>().end());
   }
 
-  /// Deletes the OpType operations in the module in the specified index.
-  static void deleteOps(ModuleOp module, int start, int end) {
+  void reduce(ModuleOp module,
+              ArrayRef<ReductionNode::Range> rangeToKeep) override {
     std::vector<Operation *> opsToRemove;
+    size_t keepIndex = 0;
 
-    for (auto op : enumerate(getSpecificOps(module))) {
+    for (auto op : enumerate(module.getOps<OpType>())) {
       int index = op.index();
-      if (index >= start && index < end)
+      if (keepIndex < rangeToKeep.size() &&
+          index == rangeToKeep[keepIndex].second)
+        ++keepIndex;
+      if (keepIndex == rangeToKeep.size() ||
+          index < rangeToKeep[keepIndex].first)
         opsToRemove.push_back(op.value());
     }
 
@@ -82,24 +69,6 @@ class OpReducer {
       o->erase();
     }
   }
-
-  /// Return the name of this reducer class.
-  StringRef getName() { return impl->getName(); }
-
-  /// Return the initial transformSpace containing the transformable indices.
-  std::vector<bool> initTransformSpace(ModuleOp module) {
-    return impl->initTransformSpace(module);
-  }
-
-  /// Generate variants by removing OpType operations from the module in the
-  /// parent and link the variants as childs in the Reduction Tree Pass.
-  void generateVariants(ReductionNode *parent, const Tester &test,
-                        int numVariants) {
-    impl->generateVariants(parent, test, numVariants, deleteOps);
-  }
-
-private:
-  std::unique_ptr<OpReducerImpl> impl;
 };
 
 } // end namespace mlir

mlir/include/mlir/Reducer/ReductionNode.h

@@ -17,82 +17,129 @@
 #ifndef MLIR_REDUCER_REDUCTIONNODE_H
 #define MLIR_REDUCER_REDUCTIONNODE_H
 
+#include <queue>
 #include <vector>
 
 #include "mlir/Reducer/Tester.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 namespace mlir {
 
-/// This class defines the ReductionNode which is used to wrap the module of
-/// a generated variant and keep track of the necessary metadata for the
-/// reduction pass. The nodes are linked together in a reduction tree structure
-/// which defines the relationship between all the different generated variants.
+/// Defines the traversal method options to be used in the reduction tree
+/// traversal.
+enum TraversalMode { SinglePath, Backtrack, MultiPath };
+
+/// This class defines the ReductionNode which is used to generate variant and
+/// keep track of the necessary metadata for the reduction pass. The nodes are
+/// linked together in a reduction tree structure which defines the relationship
+/// between all the different generated variants.
 class ReductionNode {
 public:
-  ReductionNode(ModuleOp module, ReductionNode *parent);
-
-  ReductionNode(ModuleOp module, ReductionNode *parent,
-                std::vector<bool> transformSpace);
+  template <TraversalMode mode>
+  class iterator;
 
-  /// Calculates and initializes the size and interesting values of the node.
-  void measureAndTest(const Tester &test);
+  using Range = std::pair<int, int>;
 
-  /// Returns the module.
-  ModuleOp getModule() const { return module; }
+  ReductionNode(ReductionNode *parent, std::vector<Range> range,
+                llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator);
 
-  /// Returns true if the size and interestingness have been calculated.
-  bool isEvaluated() const;
+  ReductionNode *getParent() const;
 
-  /// Returns the size in bytes of the module.
-  int getSize() const;
+  size_t getSize() const;
 
   /// Returns true if the module exhibits the interesting behavior.
-  bool isInteresting() const;
-
-  /// Returns the pointer to a child variant by index.
-  ReductionNode *getVariant(unsigned long index) const;
+  Tester::Interestingness isInteresting() const;
 
-  /// Returns the number of child variants.
-  int variantsSize() const;
+  std::vector<Range> getRanges() const;
 
-  /// Returns true if the vector containing the child variants is empty.
-  bool variantsEmpty() const;
+  std::vector<ReductionNode *> &getVariants();
 
-  /// Sort the child variants and remove the uninteresting ones.
-  void organizeVariants(const Tester &test);
+  /// Split the ranges and generate new variants.
+  std::vector<ReductionNode *> generateNewVariants();
 
-  /// Returns the number of child variants.
-  int transformSpaceSize();
-
-  /// Returns a vector indicating the transformed indices as true.
-  const std::vector<bool> getTransformSpace();
+  /// Update the interestingness result from tester.
+  void update(std::pair<Tester::Interestingness, size_t> result);
 
 private:
-  /// Link a child variant node.
-  void linkVariant(ReductionNode *newVariant);
-
-  // This is the MLIR module of this variant.
-  ModuleOp module;
-
-  // This is true if the module has been evaluated and it exhibits the
-  // interesting behavior.
-  bool interesting;
-
-  // This indicates the number of characters in the printed module if the module
-  // has been evaluated.
-  int size;
-
-  // This indicates if the module has been evaluated (measured and tested).
-  bool evaluated;
-
-  // Indicates the indices in the node that have been transformed in previous
-  // levels of the reduction tree.
-  std::vector<bool> transformSpace;
+  /// A custom BFS iterator. The difference between
+  /// llvm/ADT/BreadthFirstIterator.h is the graph we're exploring is dynamic.
+  /// We may explore more neighbors at certain node if we didn't find interested
+  /// event. As a result, we defer pushing adjacent nodes until poping the last
+  /// visited node. The graph exploration strategy will be put in
+  /// getNeighbors().
+  ///
+  /// Subclass BaseIterator and implement traversal strategy in getNeighbors().
+  template <typename T>
+  class BaseIterator {
+  public:
+    BaseIterator(ReductionNode *node) { visitQueue.push(node); }
+    BaseIterator(const BaseIterator &) = default;
+    BaseIterator() = default;
+
+    static BaseIterator end() { return BaseIterator(); }
+
+    bool operator==(const BaseIterator &i) {
+      return visitQueue == i.visitQueue;
+    }
+    bool operator!=(const BaseIterator &i) { return !(*this == i); }
+
+    BaseIterator &operator++() {
+      ReductionNode *top = visitQueue.front();
+      visitQueue.pop();
+      std::vector<ReductionNode *> neighbors = getNeighbors(top);
+      for (ReductionNode *node : neighbors)
+        visitQueue.push(node);
+      return *this;
+    }
+
+    BaseIterator operator++(int) {
+      BaseIterator tmp = *this;
+      ++*this;
+      return tmp;
+    }
+
+    ReductionNode &operator*() const { return *(visitQueue.front()); }
+    ReductionNode *operator->() const { return visitQueue.front(); }
+
+  protected:
+    std::vector<ReductionNode *> getNeighbors(ReductionNode *node) {
+      return static_cast<T *>(this)->getNeighbors(node);
+    }
+
+  private:
+    std::queue<ReductionNode *> visitQueue;
+  };
+
+  /// The size of module after applying the range constraints.
+  size_t size;
+
+  /// This is true if the module has been evaluated and it exhibits the
+  /// interesting behavior.
+  Tester::Interestingness interesting;
+
+  ReductionNode *parent;
+
+  /// We will only keep the operation with index falls into the ranges.
+  /// For example, number each function in a certain module and then we will
+  /// remove the functions with index outside the ranges and see if the
+  /// resulting module is still interesting.
+  std::vector<Range> ranges;
+
+  /// This points to the child variants that were created using this node as a
+  /// starting point.
+  std::vector<ReductionNode *> variants;
+
+  llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator;
+};
 
-  // This points to the child variants that were created using this node as a
-  // starting point.
-  std::vector<std::unique_ptr<ReductionNode>> variants;
+// Specialized iterator for SinglePath traversal
+template <>
+class ReductionNode::iterator<SinglePath>
+    : public BaseIterator<iterator<SinglePath>> {
+  friend BaseIterator<iterator<SinglePath>>;
+  using BaseIterator::BaseIterator;
+  std::vector<ReductionNode *> getNeighbors(ReductionNode *node);
 };
 
 } // end namespace mlir