Refactor the architecture of mlir-reduce

[Chia-hungDuan]

Summary

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

[Chia-hungDuan]

Add header for ReductionTreePass.cpp and fix some linting error

[rriddle]

Thanks for starting to look at this! Please fix the clang-tidy issues.

Can you also re-upload with full context?

mlir/include/mlir/Reducer/Pass

35 ↗

(On Diff #334641)

Why return vector here? Can you just return module.getOps<OpType>() directly?

48 ↗

(On Diff #334641)

Can you use ArrayRef here instead?

55 ↗

(On Diff #334641)

nit: Spell out auto here.

57 ↗

(On Diff #334641)

nit: Please drop all trivial braces in this commit.

71 ↗

(On Diff #334641)

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

29

How is this different from the other OpReducer file?

mlir/include/mlir/Reducer/ReductionTreePass.h

48–49

Does this need to be on the pass instance? It looks like it gets reset every runOnOperation.

(Same thing with the allocator, why not have it local to runOnOperation?)

mlir/tools/mlir-reduce/ReductionNode.cpp

100

nit: Use the llvm:: variants of all_of/remove_if/etc.

125

Inner comments should use //.

mlir/tools/mlir-reduce/ReductionTreePass.cpp

24

Can we just make the implementation here not templated instead? Moving to a .cpp seems like it would limit the ability for out-of-tree users to utilize this.

45

Is something erasing these clones?

[Chia-hungDuan]

Fixed the review comments and removed weird file in the previous change

[Chia-hungDuan]

mlir/include/mlir/Reducer/Pass

35 ↗

(On Diff #334641)

Just inline module.getOps<OpType>(), removed getSpecificOps

55 ↗

(On Diff #334641)

Sorry, I think this nit was required when getSepcificOps() is returning vector. I've changed to use module.getOps<OpType>(), do you think we still need to unroll it? it will be (result_pair<llvm::iterator_range<mlir::detail::op_iterator<mlir::FuncOp, mlir::Region::OpIterator>>>).

71 ↗

(On Diff #334641)

It seems that we don't need this, I removed this function. If we need this information, I think it should be involved in pass information

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

29

Sorry, the above file("mlir/include/mlir/Reducer/Pass") was added by accident. The refactor of OpReducer was supposed to be in the follow up CL. I've moved the changes into this CL

mlir/include/mlir/Reducer/ReductionTreePass.h

48–49

Moved to local vars in runOnOperation.

[Chia-hungDuan]

Reduce the number of generating new variants and change to use llvm::ArrayRef while accessing variants vector.

[jpienaar]

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

31

Could you add a comment describing the member function?

34

The previous approach here I see being generator like and towards having the same structure for all the reducers. So from the driver you'd have a common set of methods to call for each reducer. And so these now map to reduce and getNumTargetOps ?

39

It seems with this change we lose the ability to specify number of variants is that correct?

mlir/include/mlir/Reducer/ReductionNode.h

58–59

Nit: Already in private scope

68

I prefer public, protected and private ordering

mlir/include/mlir/Reducer/ReductionTreePass.h

52–53

Same re unneeded private

mlir/include/mlir/Reducer/Tester.h

38

I was thinking NotSure was like Maybe case, but it seems it is actually used to represent not tested yet, is that correct? Then perhaps Unexplored or Untested might be more intuitive

48

Check is overloaded for me (probably due to CHECK), prefer something like Return the whether file in the given path is interesting.

mlir/lib/Reducer/Tester.cpp

32

Why are you printing to temporary file here?

43

There is no bitcode here yet

mlir/tools/mlir-reduce/ReductionNode.cpp

42–44

Elide braces for trivial if bodies

44

maxElement?

96

Nit: remove this indentation, we don't have that style anywhere else here.

100

So here you are testing if any of the variants are either interesting or not? (just filtering out notsure).

mlir/tools/mlir-reduce/ReductionTreePass.cpp

26

Could TypeId be used here instead?

[Chia-hungDuan]

Add more algorithm comments and fix some lints/erros based on review comments.

mlir/lib/Reducer/Tester.cpp

32

This is used to let the test script exams the interestingness.

43

Changed to IR.

mlir/tools/mlir-reduce/ReductionNode.cpp

100

In SinglePath approach, we would like to select the most reduced variant and keep traversing from that node.
Yes, here's checking if all the variants generated from certain node are tested, if it's tested, then we can select the smallest variant, if not, we just need to wait for all the variants are tested.

I added more comments in explaining the approach.

mlir/tools/mlir-reduce/ReductionTreePass.cpp

26

We need a Op identifier but it seems that TypeID is bound to a specific value, right? If so, I think it may not be able to use in this case

[jpienaar]

I think this looks good to go and we can always iterate

mlir/tools/mlir-reduce/ReductionNode.cpp

50

Left over from debugging?

108

Why +2?

mlir/tools/mlir-reduce/ReductionTreePass.cpp

26

I was thinking one could do something like

getOpReducer(TypeId opId) {

if (opId == TypeID::get<ModuleOp>())

Perhaps I'm not following, but it seems it would give us this same structure (as we only use opType to dispatch to a particular version). But then again string comparison here is probably not on the hot path

[Chia-hungDuan]

Remove debugging header

[Chia-hungDuan]

mlir/tools/mlir-reduce/ReductionNode.cpp

50

Removed. lol

108

vector::insert inserts an element before the given iterator and it returns the iterator points to the new element. Here we insert two split ranges so the one being split now is at it + 2

mlir/tools/mlir-reduce/ReductionTreePass.cpp

26

Yes, this one only be used once while pass initialization. OTOH, we will use command line to configure the reduction process(D100155), I will revisit this structure while we have more kinds of Reducer.

[rriddle]

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

29

Can you add comments for this method? and what rangeToKeep contains?

30

nit: Drop the llvm:: here, ArrayRef is exported in the mlir namespace.

30

Should this std::pair be ReductionNode::Range?

39

Does this need to be templated? Can you pass in an OperationName to the constructor and use that for finding target ops instead?

45

Is this supposed to be in the entire module? getOps only operates on one level of the IR (i.e. in the most directly nested region).

mlir/include/mlir/Reducer/ReductionTreePass.h

53

nit: Drop the llvm:: here, StringRef is exported in the MLIR namespace.

mlir/tools/mlir-reduce/ReductionTreePass.cpp

22

LLVM/MLIR style prefers using namespace instead.

26

Inside of .cpp files, only classes should really be placed in namespaces. Functions should be marked as static and at the top level.

31

If Reducer isn't templated, we wouldn't need to abort here.

[Chia-hungDuan]

Fix linting erros.

@rriddle, about the OpReducer, I'm planing to have a registration structure for different type of Reducer in different projects. In this refactoring, I want to make the structure of graph traversal clear and extensible, I'll have another CL to address the Reducer soon. Do you think we can keep it for a short while?

Thanks

Chia-hung

[rriddle]

In D99713#2685257, @Chia-hungDuan wrote:

Fix linting erros.

@rriddle, about the OpReducer, I'm planing to have a registration structure for different type of Reducer in different projects. In this refactoring, I want to make the structure of graph traversal clear and extensible, I'll have another CL to address the Reducer soon. Do you think we can keep it for a short while?

Thanks

Chia-hung

Yeah, if you are actively cleaning that up I'm okay with deferring.

LGTM

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

 #ifndef MLIR_REDUCER_PASSES_OPREDUCER_H
 #define MLIR_REDUCER_PASSES_OPREDUCER_H
 
+#include <limits>
+
 #include "mlir/IR/Region.h"
 #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;
+  virtual void reduce(ModuleOp module,
+                      llvm::ArrayRef<std::pair<int, int>> rangeToKeep) = 0;
+  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,
+              llvm::ArrayRef<std::pair<int, int>> rangeToKeep) override {
     std::vector<Operation *> opsToRemove;
+    size_t keepIndex = 0;
 
-    for (auto op : enumerate(getSpecificOps(module))) {
+    for (auto op : llvm::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());
     }
 

       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

 #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);
-
-  /// 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;
 
   /// Returns true if the module exhibits the interesting behavior.
-  bool isInteresting() const;
+  Tester::Interestingness isInteresting() const;
 
-  /// Returns the pointer to a child variant by index.
-  ReductionNode *getVariant(unsigned long index) const;
+  std::vector<Range> getRanges() const;
 
-  /// Returns the number of child variants.
-  int variantsSize() const;
+  std::vector<ReductionNode *> getVariants() const;
 
-  /// Returns true if the vector containing the child variants is empty.
-  bool variantsEmpty() const;
+  /// Split the ranges and generate new variants.
+  std::vector<ReductionNode *> generateNewVariants();
 
-  /// Sort the child variants and remove the uninteresting ones.
-  void organizeVariants(const Tester &test);
+  /// Update the interestingness result from tester.
+  void update(std::pair<Tester::Interestingness, int> result);
 
-  /// Returns the number of child variants.
-  int transformSpaceSize();
-
-  /// Returns a vector indicating the transformed indices as true.
-  const std::vector<bool> getTransformSpace();
+private:
+  /// Split each of the range into half. Return false if all ranges have length
+  /// 1, i.e., we can't do further split.
+  bool splitRanges();
 
 private:
-  /// Link a child variant node.
-  void linkVariant(ReductionNode *newVariant);
+  /// 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.
+  ///
+  /// Subclass BaseIterator and implement your node exploration strategy inside
+  /// getNeighbors().
+  template <typename T>
+  class BaseIterator {
+  protected:
+    std::vector<ReductionNode *> getNeighbors(ReductionNode *node) {
+      return static_cast<T *>(this)->getNeighbors(node);
+    }
+
+  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(); }
+
+  private:
+    std::queue<ReductionNode *> visitQueue;
+  };
 
-  // This is the MLIR module of this variant.
-  ModuleOp module;
+public:
+  template <TraversalMode mode>
+  class iterator;
 
-  // This is true if the module has been evaluated and it exhibits the
-  // interesting behavior.
-  bool interesting;
+  template <>
+  class iterator<SinglePath> : public BaseIterator<iterator<SinglePath>> {
+    friend BaseIterator<iterator<SinglePath>>;
+    using BaseIterator::BaseIterator;
+    std::vector<ReductionNode *> getNeighbors(ReductionNode *node);
+  };
 
-  // This indicates the number of characters in the printed module if the module
-  // has been evaluated.
+private:
+  /// The size of module after applying the range constraints.
   int size;
 
-  // This indicates if the module has been evaluated (measured and tested).
-  bool evaluated;
+  /// 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;
 
-  // Indicates the indices in the node that have been transformed in previous
-  // levels of the reduction tree.
-  std::vector<bool> transformSpace;
+  /// This points to the child variants that were created using this node as a
+  /// starting point.
+  std::vector<ReductionNode *> variants;
 
-  // This points to the child variants that were created using this node as a
-  // starting point.
-  std::vector<std::unique_ptr<ReductionNode>> variants;
+  llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator;
 };
 
 } // end namespace mlir

mlir/include/mlir/Reducer/ReductionTreePass.h

 #include "PassDetail.h"
 #include "ReductionNode.h"
 #include "mlir/Reducer/Passes/OpReducer.h"
-#include "mlir/Reducer/ReductionTreeUtils.h"
 #include "mlir/Reducer/Tester.h"
 
 #define DEBUG_TYPE "mlir-reduce"
 
 namespace mlir {
 
-// Defines the traversal method options to be used in the reduction tree
-/// traversal.
-enum TraversalMode { SinglePath, Backtrack, MultiPath };
-
 /// This class defines the Reduction Tree Pass. It provides a framework to
 /// to implement a reduction pass using a tree structure to keep track of the
 /// generated reduced variants.
-template <typename Reducer, TraversalMode mode>
-class ReductionTreePass
-    : public ReductionTreeBase<ReductionTreePass<Reducer, mode>> {
+class ReductionTreePass : public ReductionTreeBase<ReductionTreePass> {
 public:
   ReductionTreePass(const ReductionTreePass &pass)
-      : ReductionTreeBase<ReductionTreePass<Reducer, mode>>(pass),
-        root(new ReductionNode(pass.root->getModule().clone(), nullptr)),
-        test(pass.test) {}
+      : ReductionTreeBase<ReductionTreePass>(pass), opType(pass.opType),
+        mode(pass.mode), test(pass.test) {}
 
-  ReductionTreePass(const Tester &test) : test(test) {}
+  ReductionTreePass(llvm::StringRef opType, TraversalMode mode,
+                    const Tester &test)
+      : opType(opType), mode(mode), test(test) {}
 
   /// Runs the pass instance in the pass pipeline.
-  void runOnOperation() override {
-    ModuleOp module = this->getOperation();
-    Reducer reducer;
-    std::vector<bool> transformSpace = reducer.initTransformSpace(module);
-    ReductionNode *reduced;
-
-    this->root =
-        std::make_unique<ReductionNode>(module, nullptr, transformSpace);
-
-    root->measureAndTest(test);
+  void runOnOperation() override;
 
-    LLVM_DEBUG(llvm::dbgs() << "\nReduction Tree Pass: " << reducer.getName(););
-    switch (mode) {
-    case SinglePath:
-      LLVM_DEBUG(llvm::dbgs() << " (Single Path)\n";);
-      reduced = singlePathTraversal();
-      break;
-    default:
-      llvm::report_fatal_error("Traversal method not currently supported.");
-    }
-
-    ReductionTreeUtils::updateGoldenModule(module,
-                                           reduced->getModule().clone());
-  }
+private:
+  template <TraversalMode mode>
+  ModuleOp findOptimal(ModuleOp module, std::unique_ptr<OpReducer> reducer,
+                       ReductionNode *node);
 
 private:
-  // Points to the root node in this reduction tree.
-  std::unique_ptr<ReductionNode> root;
+  /// The name of operation that we will try to remove.
+  llvm::StringRef opType;
 
-  // This object defines the variant generation at each level of the reduction
-  // tree.
-  Reducer reducer;
+  TraversalMode mode;
 
-  // This is used to test the interesting behavior of the reduction nodes in the
-  // tree.
+  /// This is used to test the interesting behavior of the reduction nodes in
+  /// the tree.
   const Tester &test;
-
-  /// Traverse the most reduced path in the reduction tree by generating the
-  /// variants at each level using the Reducer parameter's generateVariants
-  /// function. Stops when no new successful variants can be created at the
-  /// current level.
-  ReductionNode *singlePathTraversal() {
-    ReductionNode *currNode = root.get();
-    ReductionNode *smallestNode = currNode;
-    int tSpaceSize = currNode->transformSpaceSize();
-    std::vector<int> path;
-
-    ReductionTreeUtils::updateSmallestNode(currNode, smallestNode, path);
-
-    LLVM_DEBUG(llvm::dbgs() << "\nGenerating 1 variant: applying the ");
-    LLVM_DEBUG(llvm::dbgs() << "transformation to the entire module\n");
-
-    reducer.generateVariants(currNode, test, 1);
-    LLVM_DEBUG(llvm::dbgs() << "Testing\n");
-    currNode->organizeVariants(test);
-
-    if (!currNode->variantsEmpty())
-      return currNode->getVariant(0);
-
-    while (tSpaceSize != 1) {
-      ReductionTreeUtils::updateSmallestNode(currNode, smallestNode, path);
-
-      LLVM_DEBUG(llvm::dbgs() << "\nGenerating 2 variants: applying the ");
-      LLVM_DEBUG(llvm::dbgs() << "transformation to two different sections ");
-      LLVM_DEBUG(llvm::dbgs() << "of transformable indices\n");
-
-      reducer.generateVariants(currNode, test, 2);
-      LLVM_DEBUG(llvm::dbgs() << "Testing\n");
-      currNode->organizeVariants(test);
-
-      if (currNode->variantsEmpty())
-        break;
-
-      currNode = currNode->getVariant(0);
-      tSpaceSize = currNode->transformSpaceSize();
-      path.push_back(0);
-    }
-
-    if (tSpaceSize == 1) {
-      ReductionTreeUtils::updateSmallestNode(currNode, smallestNode, path);
-
-      LLVM_DEBUG(llvm::dbgs() << "\nGenerating 1 variants: applying the ");
-      LLVM_DEBUG(llvm::dbgs() << "transformation to the only transformable");
-      LLVM_DEBUG(llvm::dbgs() << "index\n");
-
-      reducer.generateVariants(currNode, test, 1);
-      LLVM_DEBUG(llvm::dbgs() << "Testing\n");
-      currNode->organizeVariants(test);
-
-      if (!currNode->variantsEmpty()) {
-        currNode = currNode->getVariant(0);
-        path.push_back(0);
-
-        ReductionTreeUtils::updateSmallestNode(currNode, smallestNode, path);
-      }
-    }
-
-    return currNode;
-  }
 };
 
 } // end namespace mlir

mlir/include/mlir/Reducer/ReductionTreeUtils.h

-//===- ReductionTreeUtils.h - Reduction Tree utilities ----------*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file defines the Reduction Tree Utilities. It defines pass independent
-// methods that help in the reduction passes of the MLIR Reduce tool.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_REDUCER_REDUCTIONTREEUTILS_H
-#define MLIR_REDUCER_REDUCTIONTREEUTILS_H
-
-#include <tuple>
-
-#include "PassDetail.h"
-#include "ReductionNode.h"
-#include "mlir/Reducer/Tester.h"
-#include "llvm/Support/Debug.h"
-
-namespace mlir {
-
-// Defines the utilities for the implementation of custom reduction
-// passes using the ReductionTreePass framework.
-namespace ReductionTreeUtils {
-
-/// Update the golden module's content with that of the reduced module.
-void updateGoldenModule(ModuleOp &golden, ModuleOp reduced);
-
-/// Update the smallest node traversed so far in the reduction tree and
-/// print the debugging information for the currNode being traversed.
-void updateSmallestNode(ReductionNode *currNode, ReductionNode *&smallestNode,
-                        std::vector<int> path);
-
-/// Create a transform space index vector based on the specified number of
-/// indices.
-std::vector<bool> createTransformSpace(ModuleOp module, int numIndices);
-
-/// Create the specified number of variants by applying the transform method
-/// to different ranges of indices in the parent module. The isDeletion boolean
-/// specifies if the transformation is the deletion of indices.
-void createVariants(ReductionNode *parent, const Tester &test, int numVariants,
-                    llvm::function_ref<void(ModuleOp, int, int)> transform,
-                    bool isDeletion);
-
-} // namespace ReductionTreeUtils
-
-} // end namespace mlir
-
-#endif

mlir/include/mlir/Reducer/Tester.h

 /// case file.
 class Tester {
 public:
+  enum class Interestingness {
+    True,
+    False,
+    NotSure,
+  };
+
   Tester(StringRef testScript, ArrayRef<std::string> testScriptArgs);
 
   /// Runs the interestingness testing script on a MLIR test case file. Returns
   /// true if the interesting behavior is present in the test case or false
   /// otherwise.
-  bool isInteresting(StringRef testCase) const;
+  std::pair<Interestingness, int> isInteresting(ModuleOp module) const;
+
+  /// Check the interestingness of the file in the given path.
+  Interestingness isInteresting(StringRef testCase) const;
 
 private:
   StringRef testScript;

mlir/lib/Reducer/Tester.cpp

 
 #include "mlir/Reducer/Tester.h"
 
+#include "llvm/Support/ToolOutputFile.h"
+
 using namespace mlir;
 
 Tester::Tester(StringRef scriptName, ArrayRef<std::string> scriptArgs)
     : testScript(scriptName), testScriptArgs(scriptArgs) {}
 
+std::pair<Tester::Interestingness, int>
+Tester::isInteresting(ModuleOp module) const {
+  SmallString<128> filepath;
+  int fd;
+
+  // Print module to temporary file.
+  std::error_code ec =
+      llvm::sys::fs::createTemporaryFile("mlir-reduce", "mlir", fd, filepath);
+
+  if (ec)
+    llvm::report_fatal_error("Error making unique filename: " + ec.message());
+
+  llvm::ToolOutputFile out(filepath, fd);
+  module.print(out.os());
+  out.os().close();
+
+  if (out.os().has_error())
+    llvm::report_fatal_error("Error emitting bitcode to file '" + filepath);
+
+  int size = out.os().tell();
+  return std::make_pair(isInteresting(filepath), size);
+}
+
 /// Runs the interestingness testing script on a MLIR test case file. Returns
 /// true if the interesting behavior is present in the test case or false
 /// otherwise.
-bool Tester::isInteresting(StringRef testCase) const {
+Tester::Interestingness Tester::isInteresting(StringRef testCase) const {
 
   std::vector<StringRef> testerArgs;
   testerArgs.push_back(testCase);

                              false);
 
   if (!result)
-    return false;
+    return Interestingness::False;
 
-  return true;
+  return Interestingness::True;
 }

mlir/tools/mlir-reduce/CMakeLists.txt

 
 add_llvm_tool(mlir-reduce
   OptReductionPass.cpp
-  Passes/OpReducer.cpp
   ReductionNode.cpp
-  ReductionTreeUtils.cpp
+  ReductionTreePass.cpp
   mlir-reduce.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/tools/mlir-reduce/OptReductionPass.cpp

   PassManager pmTransform(context);
   pmTransform.addPass(std::move(optPass));
 
+  std::pair<Tester::Interestingness, int> original = test.isInteresting(module);
+
   if (failed(pmTransform.run(moduleVariant)))
     return;
 
-  ReductionNode original(module, nullptr);
-  original.measureAndTest(test);
-
-  ReductionNode reduced(moduleVariant, nullptr);
-  reduced.measureAndTest(test);
+  std::pair<Tester::Interestingness, int> reduced =
+      test.isInteresting(moduleVariant);
 
-  if (reduced.isInteresting() && reduced.getSize() < original.getSize()) {
-    ReductionTreeUtils::updateGoldenModule(module, reduced.getModule().clone());
+  if (reduced.first == Tester::Interestingness::True &&
+      reduced.second < original.second) {
+    module.getBody()->clear();
+    module.getBody()->getOperations().splice(
+        module.getBody()->begin(), moduleVariant.getBody()->getOperations());
     LLVM_DEBUG(llvm::dbgs() << "\nSuccessful Transformed version\n\n");
   } else {
     LLVM_DEBUG(llvm::dbgs() << "\nUnsuccessful Transformed version\n\n");
   }
 
+  moduleVariant->destroy();
+
   LLVM_DEBUG(llvm::dbgs() << "Pass Complete\n\n");
 }

mlir/tools/mlir-reduce/Passes/OpReducer.cpp

-//===- OpReducer.cpp - Operation Reducer ------------------------*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file defines the OpReducer class. It defines a variant generator method
-// with the purpose of producing different variants by eliminating a
-// parameterizable type of operations from the  parent module.
-//
-//===----------------------------------------------------------------------===//
-#include "mlir/Reducer/Passes/OpReducer.h"
-
-using namespace mlir;
-
-OpReducerImpl::OpReducerImpl(
-    llvm::function_ref<std::vector<Operation *>(ModuleOp)> getSpecificOps)
-    : getSpecificOps(getSpecificOps) {}
-
-/// Return the name of this reducer class.
-StringRef OpReducerImpl::getName() {
-  return StringRef("High Level Operation Reduction");
-}
-
-/// Return the initial transformSpace containing the transformable indices.
-std::vector<bool> OpReducerImpl::initTransformSpace(ModuleOp module) {
-  auto ops = getSpecificOps(module);
-  int numOps = std::distance(ops.begin(), ops.end());
-  return ReductionTreeUtils::createTransformSpace(module, numOps);
-}
-
-/// Generate variants by removing opType operations from the module in the
-/// parent and link the variants as childs in the Reduction Tree Pass.
-void OpReducerImpl::generateVariants(
-    ReductionNode *parent, const Tester &test, int numVariants,
-    llvm::function_ref<void(ModuleOp, int, int)> transform) {
-  ReductionTreeUtils::createVariants(parent, test, numVariants, transform,
-                                     true);
-}

mlir/tools/mlir-reduce/ReductionNode.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Reducer/ReductionNode.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <algorithm>
 
 using namespace mlir;
 
-/// Sets up the metadata and links the node to its parent.
-ReductionNode::ReductionNode(ModuleOp module, ReductionNode *parent)
-    : module(module), evaluated(false) {
+ReductionNode::ReductionNode(
+    ReductionNode *parent, std::vector<Range> ranges,
+    llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator)
+    : interesting(Tester::Interestingness::NotSure),
+      /// Root node will have the parent pointer point to themselves.
+      parent(parent == nullptr ? this : parent), ranges(ranges),
+      allocator(allocator) {}
 
-  if (parent != nullptr)
-    parent->linkVariant(this);
-}
+/// Returns the size in bytes of the module.
+int ReductionNode::getSize() const { return size; }
 
-ReductionNode::ReductionNode(ModuleOp module, ReductionNode *parent,
-                             std::vector<bool> transformSpace)
-    : module(module), evaluated(false), transformSpace(transformSpace) {
+ReductionNode *ReductionNode::getParent() const { return parent; }
 
-  if (parent != nullptr)
-    parent->linkVariant(this);
+/// Returns true if the module exhibits the interesting behavior.
+Tester::Interestingness ReductionNode::isInteresting() const {
+  return interesting;
 }
 
-/// Calculates and updates the size and interesting values of the module.
-void ReductionNode::measureAndTest(const Tester &test) {
-  SmallString<128> filepath;
-  int fd;
-
-  // Print module to temporary file.
-  std::error_code ec =
-      llvm::sys::fs::createTemporaryFile("mlir-reduce", "mlir", fd, filepath);
-
-  if (ec)
-    llvm::report_fatal_error("Error making unique filename: " + ec.message());
-
-  llvm::ToolOutputFile out(filepath, fd);
-  module.print(out.os());
-  out.os().close();
-
-  if (out.os().has_error())
-    llvm::report_fatal_error("Error emitting bitcode to file '" + filepath);
+std::vector<ReductionNode::Range> ReductionNode::getRanges() const {
+  return ranges;
+}
 
-  size = out.os().tell();
-  interesting = test.isInteresting(filepath);
-  evaluated = true;
+bool ReductionNode::splitRanges() {
+  std::vector<Range> newRange;
+  bool split = false;
+  for (const Range &range : ranges) {
+    if (range.second - range.first <= 1) {
+      newRange.push_back(range);
+    } else {
+      split = true;
+      int half = (range.first + range.second) / 2;
+      newRange.push_back(std::make_pair(range.first, half));
+      newRange.push_back(std::make_pair(half, range.second));
+    }
+  }
+  if (split) {
+    ranges = newRange;
+  }
+  return split;
 }
 
-/// Returns true if the size and interestingness have been calculated.
-bool ReductionNode::isEvaluated() const { return evaluated; }
+std::vector<ReductionNode *> ReductionNode::getVariants() const {
+  return variants;
+}
 
-/// Returns the size in bytes of the module.
-int ReductionNode::getSize() const { return size; }
+std::vector<ReductionNode *> ReductionNode::generateNewVariants() {
+  // The ranges doesn't change, no new variant will be generated.
+  if (!splitRanges()) {
+    return {};
+  }
 
-/// Returns true if the module exhibits the interesting behavior.
-bool ReductionNode::isInteresting() const { return interesting; }
+  std::vector<ReductionNode *> newNodes;
 
-/// Returns the pointers to the child variants.
-ReductionNode *ReductionNode::getVariant(unsigned long index) const {
-  if (index < variants.size())
-    return variants[index].get();
+  for (const Range &range : ranges) {
+    std::vector<Range> subRange = ranges;
+    llvm::erase_value(subRange, range);
+    ReductionNode *newNode = allocator.Allocate();
+    new (newNode) ReductionNode(this, subRange, allocator);
+    newNodes.push_back(newNode);
+    variants.push_back(newNode);
+  }
 
-  return nullptr;
+  return newNodes;
 }
 
-/// Returns the number of child variants.
-int ReductionNode::variantsSize() const { return variants.size(); }
-
-/// Returns true if the child variants vector is empty.
-bool ReductionNode::variantsEmpty() const { return variants.empty(); }
-
-/// Link a child variant node.
-void ReductionNode::linkVariant(ReductionNode *newVariant) {
-  std::unique_ptr<ReductionNode> ptrVariant(newVariant);
-  variants.push_back(std::move(ptrVariant));
+void ReductionNode::update(std::pair<Tester::Interestingness, int> result) {
+  std::tie(interesting, size) = result;
 }
 
-/// Sort the child variants and remove the uninteresting ones.
-void ReductionNode::organizeVariants(const Tester &test) {
-  // Ensure all variants are evaluated.
-  for (auto &var : variants)
-    if (!var->isEvaluated())
-      var->measureAndTest(test);
-
-  // Sort variants by interestingness and size.
-  llvm::array_pod_sort(
-      variants.begin(), variants.end(), [](const auto *lhs, const auto *rhs) {
-        if (lhs->get()->isInteresting() && !rhs->get()->isInteresting())
-          return 0;
-
-        if (!lhs->get()->isInteresting() && rhs->get()->isInteresting())
-          return 1;
-
-        return (lhs->get()->getSize(), rhs->get()->getSize());
-      });
-
-  int interestingCount = 0;
-  for (auto &var : variants) {
-    if (var->isInteresting()) {
-      ++interestingCount;
-    } else {
-      break;
-    }
+std::vector<ReductionNode *>
+ReductionNode::iterator<SinglePath>::getNeighbors(ReductionNode *node) {
+  // Single Path: Traverses the smallest successful variant at each level until
+  //              no new successful variants can be created at that level.
+  std::vector<ReductionNode *> variantsFromParent =
+      node->getParent()->getVariants();
+
+  if (!llvm::all_of(variantsFromParent, [](ReductionNode *node) {
+        return node->isInteresting() != Tester::Interestingness::NotSure;
+      })) {
+    // We still have variants in queue, wait for them to test interestingness
+    // and then we can try to select the smallest variant.
+    return {};
   }
 
-  // Remove uninteresting variants.
-  variants.resize(interestingCount);
-}
-
-/// Returns the number of non transformed indices.
-int ReductionNode::transformSpaceSize() {
-  return std::count(transformSpace.begin(), transformSpace.end(), false);
-}
+  variantsFromParent.erase(
+      llvm::remove_if(variantsFromParent,
+                      [](ReductionNode *node) {
+                        return node->isInteresting() ==
+                               Tester::Interestingness::False;
+                      }),
+      variantsFromParent.end());
+
+  auto smallest =
+      std::min_element(variantsFromParent.begin(), variantsFromParent.end(),
+                       [](ReductionNode *lhs, ReductionNode *rhs) {
+                         return lhs->getSize() < rhs->getSize();
+                       });
+
+  if (smallest != variantsFromParent.end()) {
+    // We got a smallest one, keep traversing from this node.
+    node = *smallest;
+  } else {
+    // None of these variants is interesting, let the parent node to generate
+    // more variants.
+    node = node->getParent();
+  }
 
-/// Returns a vector of the transformable indices in the Module.
-const std::vector<bool> ReductionNode::getTransformSpace() {
-  return transformSpace;
+  return node->generateNewVariants();
 }

mlir/tools/mlir-reduce/ReductionTreePass.cpp

+//===- ReductionTreePass.cpp - ReductionTreePass Implementation -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the Reduction Tree Pass class. It provides a framework for
+// the implementation of different reduction passes in the MLIR Reduce tool. It
+// allows for custom specification of the variant generation behavior. It
+// implements methods that define the different possible traversals of the
+// reduction tree.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Reducer/ReductionTreePass.h"
+
+#include "llvm/Support/Allocator.h"
+
+namespace mlir {
+
+namespace {
+
+std::unique_ptr<OpReducer> getOpReducer(llvm::StringRef opType) {
+  if (opType == ModuleOp::getOperationName())
+    return std::make_unique<Reducer<ModuleOp>>();
+  else if (opType == FuncOp::getOperationName())
+    return std::make_unique<Reducer<FuncOp>>();
+  llvm_unreachable("Now only supports two built-in ops");
+}
+
+} // end namespace
+
+void ReductionTreePass::runOnOperation() {
+  ModuleOp module = this->getOperation();
+  std::unique_ptr<OpReducer> reducer = getOpReducer(opType);
+  std::vector<std::pair<int, int>> ranges = {
+      {0, reducer->getNumTargetOps(module)}};
+
+  llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;
+
+  ReductionNode *root = allocator.Allocate();
+  new (root) ReductionNode(nullptr, ranges, allocator);
+
+  ModuleOp golden = module;
+  switch (mode) {
+  case TraversalMode::SinglePath:
+    golden = findOptimal<SinglePath>(module, std::move(reducer), root);
+    break;
+  default:
+    llvm_unreachable("Unsupported mode");
+  }
+
+  if (golden != module) {
+    module.getBody()->clear();
+    module.getBody()->getOperations().splice(module.getBody()->begin(),
+                                             golden.getBody()->getOperations());
+    golden->destroy();
+  }
+}
+
+template <TraversalMode mode>
+ModuleOp ReductionTreePass::findOptimal(ModuleOp module,
+                                        std::unique_ptr<OpReducer> reducer,
+                                        ReductionNode *root) {
+  std::pair<Tester::Interestingness, int> initStatus =
+      test.isInteresting(module);
+  root->update(initStatus);
+
+  ReductionNode *smallestNode = root;
+  ModuleOp golden = module;
+
+  ReductionNode::iterator<mode> iter(root);
+
+  while (iter != ReductionNode::iterator<mode>::end()) {
+    ModuleOp cloneModule = module.clone();
+
+    ReductionNode &currentNode = *iter;
+    reducer->reduce(cloneModule, currentNode.getRanges());
+
+    std::pair<Tester::Interestingness, int> result =
+        test.isInteresting(cloneModule);
+    currentNode.update(result);
+
+    if (result.first == Tester::Interestingness::True &&
+        result.second < smallestNode->getSize()) {
+      smallestNode = &currentNode;
+      golden = cloneModule;
+    } else {
+      cloneModule->destroy();
+    }
+
+    ++iter;
+  }
+
+  return golden;
+}
+
+} // end namespace mlir

mlir/tools/mlir-reduce/ReductionTreeUtils.cpp

-//===- ReductionTreeUtils.cpp - Reduction Tree Utilities ------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file defines the Reduction Tree Utilities. It defines pass independent
-// methods that help in a reduction pass of the MLIR Reduce tool.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Reducer/ReductionTreeUtils.h"
-
-#define DEBUG_TYPE "mlir-reduce"
-
-using namespace mlir;
-
-/// Update the golden module's content with that of the reduced module.
-void ReductionTreeUtils::updateGoldenModule(ModuleOp &golden,
-                                            ModuleOp reduced) {
-  golden.getBody()->clear();
-
-  golden.getBody()->getOperations().splice(golden.getBody()->begin(),
-                                           reduced.getBody()->getOperations());
-}
-
-/// Update the smallest node traversed so far in the reduction tree and
-/// print the debugging information for the currNode being traversed.
-void ReductionTreeUtils::updateSmallestNode(ReductionNode *currNode,
-                                            ReductionNode *&smallestNode,
-                                            std::vector<int> path) {
-  LLVM_DEBUG(llvm::dbgs() << "\nTree Path: root");
-  #ifndef NDEBUG
-  for (int nodeIndex : path)
-    LLVM_DEBUG(llvm::dbgs() << " -> " << nodeIndex);
-  #endif
-
-  LLVM_DEBUG(llvm::dbgs() << "\nSize (chars): " << currNode->getSize());
-  if (currNode->getSize() < smallestNode->getSize()) {
-    LLVM_DEBUG(llvm::dbgs() << " - new smallest node!");
-    smallestNode = currNode;
-  }
-}
-
-/// Create a transform space index vector based on the specified number of
-/// indices.
-std::vector<bool> ReductionTreeUtils::createTransformSpace(ModuleOp module,
-                                                           int numIndices) {
-  std::vector<bool> transformSpace;
-  for (int i = 0; i < numIndices; ++i)
-    transformSpace.push_back(false);
-
-  return transformSpace;
-}
-
-/// Translate section start and end into a vector of ranges specifying the
-/// section in the non transformed indices in the transform space.
-static std::vector<std::tuple<int, int>> getRanges(std::vector<bool> tSpace,
-                                                   int start, int end) {
-  std::vector<std::tuple<int, int>> ranges;
-  int rangeStart = 0;
-  int rangeEnd = 0;
-  bool inside = false;
-  int transformableCount = 0;
-
-  for (auto element : llvm::enumerate(tSpace)) {
-    int index = element.index();
-    bool value = element.value();
-
-    if (start <= transformableCount && transformableCount < end) {
-      if (!value && !inside) {
-        inside = true;
-        rangeStart = index;
-      }
-      if (value && inside) {
-        rangeEnd = index;
-        ranges.push_back(std::make_tuple(rangeStart, rangeEnd));
-        inside = false;
-      }
-    }
-
-    if (!value)
-      transformableCount++;
-
-    if (transformableCount == end && inside) {
-      ranges.push_back(std::make_tuple(rangeStart, index + 1));
-      inside = false;
-      break;
-    }
-  }
-
-  return ranges;
-}
-
-/// Create the specified number of variants by applying the transform method
-/// to different ranges of indices in the parent module. The isDeletion boolean
-/// specifies if the transformation is the deletion of indices.
-void ReductionTreeUtils::createVariants(
-    ReductionNode *parent, const Tester &test, int numVariants,
-    llvm::function_ref<void(ModuleOp, int, int)> transform, bool isDeletion) {
-  std::vector<bool> newTSpace;
-  ModuleOp module = parent->getModule();
-
-  std::vector<bool> parentTSpace = parent->getTransformSpace();
-  int indexCount = parent->transformSpaceSize();
-  std::vector<std::tuple<int, int>> ranges;
-
-  // No new variants can be created.
-  if (indexCount == 0)
-    return;
-
-  // Create a single variant by transforming the unique index.
-  if (indexCount == 1) {
-    ModuleOp variantModule = module.clone();
-    if (isDeletion) {
-      transform(variantModule, 0, 1);
-    } else {
-      ranges = getRanges(parentTSpace, 0, parentTSpace.size());
-      transform(variantModule, std::get<0>(ranges[0]), std::get<1>(ranges[0]));
-    }
-
-    new ReductionNode(variantModule, parent, newTSpace);
-
-    return;
-  }
-
-  // Create the specified number of variants.
-  for (int i = 0; i < numVariants; ++i) {
-    ModuleOp variantModule = module.clone();
-    newTSpace = parent->getTransformSpace();
-    int sectionSize = indexCount / numVariants;
-    int sectionStart = sectionSize * i;
-    int sectionEnd = sectionSize * (i + 1);
-
-    if (i == numVariants - 1)
-      sectionEnd = indexCount;
-
-    if (isDeletion)
-      transform(variantModule, sectionStart, sectionEnd);
-
-    ranges = getRanges(parentTSpace, sectionStart, sectionEnd);
-
-    for (auto range : ranges) {
-      int rangeStart = std::get<0>(range);
-      int rangeEnd = std::get<1>(range);
-
-      for (int x = rangeStart; x < rangeEnd; ++x)
-        newTSpace[x] = true;
-
-      if (!isDeletion)
-        transform(variantModule, rangeStart, rangeEnd);
-    }
-
-    // Create Reduction Node in the Reduction tree
-    new ReductionNode(variantModule, parent, newTSpace);
-  }
-}

mlir/tools/mlir-reduce/mlir-reduce.cpp

   // Initialize test environment.
   const Tester test(testFilename, testArguments);
 
-  if (!test.isInteresting(inputFilename))
+  if (test.isInteresting(inputFilename) != Tester::Interestingness::True)
     llvm::report_fatal_error(
         "Input test case does not exhibit interesting behavior");
 

 
   } else if (passTestSpecifier == "function-reducer") {
 
-    // Reduction tree pass with OpReducer variant generation and single path
+    // Reduction tree pass with Reducer variant generation and single path
     // traversal.
-    pm.addPass(
-        std::make_unique<ReductionTreePass<OpReducer<FuncOp>, SinglePath>>(
-            test));
+    pm.addPass(std::make_unique<ReductionTreePass>(
+        FuncOp::getOperationName(), TraversalMode::SinglePath, test));
   }
 
   ModuleOp m = moduleRef.get().clone();