[mlir-reduce] Reducer refactor.

[Chia-hungDuan]

Summary

• A Reducer is a kind of RewritePattern, so it's just the same as

writing graph rewrite.

• ReductionTreePass operates on Operation rather than ModuleOp, so that
• we are able to reduce a nested structure(e.g., module in module) by
• self-nesting.

[Chia-hungDuan]

Handle empty region and adjusted some early return logics.

[Chia-hungDuan]

Fixed linting error and updated comment

[rriddle]

Thanks! Took a first pass.

mlir/include/mlir/Reducer/ReducerRegistry.h

1 ↗

(On Diff #339968)

18 ↗

(On Diff #339968)

Can you document these?

mlir/include/mlir/Reducer/ReductionNode.h

65

Can you add documentation here?

71

Should you be using ArrayRef here instead?

94

Please document this.

mlir/lib/Reducer/ReductionNode.cpp

62–64

67–68

Given you are already constructing the vector, can you move it when creating the new node? To avoid the extra copy.

mlir/lib/Reducer/ReductionTreePass.cpp

31–32

We should generally try to avoid static variables whenever reasonably possible. Why does this need to be static?

39

Why can't you use the default constructor? Pass options are copied implicitly, so I wouldn't expect you to need this.

50

nit: Drop trivial braces here and below.

52–53

Can you move this (pattern construction) to the initialize method?

https://mlir.llvm.org/docs/PassManagement/#initialization

62–63

You would also need to ensure that the nested operation is isolated from above, otherwise it will assert.

64–66

Can you cache these? This could end up creating a lot of duplicate pass managers, whereas you should only need one of each op-type.

66

clone()?

85–86

Does this need to be on ReductionTreePass? Can you make this static to this file instead? Also, can you split up this function? It is quite large.

119

Drop the const here, ArrayRef has only const API.

183–187

Could this be pass failures instead of hard crashes?

mlir/test/lib/Reducer/MLIRTestReducer.cpp

41

[Chia-hungDuan]

mlir/include/mlir/Reducer/ReductionNode.h

65

Added the comment at the data member section.

mlir/lib/Reducer/ReductionTreePass.cpp

31–32

This is supposed to have the same function as passRegistry to let the users register various kinds of reducer patterns. The size will be far smaller than passRegistry in fact.

39

Done.
I thought copy constructor would do that, but found that Option are not copied. I forgot there's a clone() and it handles the Option properly.

---

std::unique_ptr<Pass> clone() const {
  auto newInst = clonePass();
  newInst->copyOptionValuesFrom(this);
  return newInst;
}

---

Do you think we need to add this in the copy ctor of Pass?

62–63

Thanks for reminder, I didn't notice that

64–66

Rewrite it without using OpPassManager, just do a operation walk recursively.

85–86

Split the lambda to a static function and clean the content in this function. Let's see if we still need to move it to static in next review.

183–187

These two are supposed to be always true. So if it's failed then means we have problems in the algorithm. Or do you think using an assertion will be better?

[Chia-hungDuan]

Add comments and did some structures clean up

[rriddle]

mlir/include/mlir/Reducer/ReductionNode.h

65

Can you add it here too? Data members aren't necessarily user facing API.

71

Drop the llvm::, ArrayRef is exported in the mlir namespace. Same applies throughout the revision,

mlir/lib/Reducer/ReductionNode.cpp

39–50

nit: You could just inline these into the header.

40

67–68

Using ArrayRef will still generate a copy?

71

Can you use getVariants().drop_begin?

mlir/lib/Reducer/ReductionTreePass.cpp

31–32

I don't understand your point, why can't the user just collect the registrations and pass them when constructing the pass? The pass registration mechanism is a bit of a special case (and could likely be restructured to not be static, albeit with a bit of work).

52

Drop the trivial braces here.

57

nit: Fix the naming style here.

93

nit: Fix the naming style here.

114

Can you make this a non-member of the class?

183–187

LLVM style is to assert liberally for invariants. If these things generally shouldn't happen, assert is fine.

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

90

If we generate something invalid, couldn't we prevent that from escaping the pass boundary? Do you have an example use case that requires this?

[Chia-hungDuan]

mlir/lib/Reducer/ReductionNode.cpp

67–68

Ah, my bad. Did few revisions here and forgot the copy issue. Fixed

mlir/lib/Reducer/ReductionTreePass.cpp

31–32

I didn't find a better way to pass the patterns to construct a Pass.

Now I added a function ReductionTreePass::populateReducerPatterns to collect the patterns. It'll iterate over the pass manager and pass the patterns if it's a ReductionTreePass.

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

90

• The original ---

func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {

cond_br %arg0, ^bb1, ^bb2

^bb1:

br ^bb3(%arg1 : memref<2xf32>)

^bb2:

%0 = memref.alloc() : memref<2xf32>                                               
br ^bb3(%0 : memref<2xf32>)

^bb3(%1: memref<2xf32>):

return

}

• reduced module ---

module {

func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
  cond_br %arg0, ^bb1, ^bb2
^bb1:  // pred: ^bb0
^bb2:  // pred: ^bb0
^bb3(%0: memref<2xf32>):  // no predecessors
}

}

• The error message ---

error: empty block: expect at least a terminator

This is the test case 'test/mlir-reduce/simple-test.mlir', the tester always return true so that we reduce it into this invalid form.

Yes, I'm thinking if we should always check the validity of the reduced form. Now I let the tester to decide if an invalid form is acceptable. What do you think?

[Chia-hungDuan]

Addressed review comments

[jpienaar]

Nice, first scan

mlir/include/mlir/Reducer/ReductionNode.h

38–42

ReductionTreePass

78

Lets be consistent here use "///" and comment all members.

161–168

represents the selected subset?

[jpienaar]

mlir/include/mlir/Reducer/ReductionNode.h

171

reduced

171

I think you have all the correct info here, perhaps we could make it shorter though. WDYT of

startRanges records the ranges of operations selected from the parent node to produce this ReductionNode. It can be used to construct the reduction path from the root.

(one could include the i.e., section still).

mlir/lib/Reducer/ReductionNode.cpp

64–66

Would getRange().size() == 0 be possible here? (it would probably not produce something post, I was just curios about != 1 vs > 1 here)

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

90

I think that is a good idea: the input to test on should be valid, else we may be reducing a case where a pass produces an invalid IR, and we end up retaining the error message by creating a new invalid input with the same failure message instead.

102

You should be able to drop llvm:: here and drop the ReductionTreePass on lhs as we have it captured in the cast

[Chia-hungDuan]

mlir/include/mlir/Reducer/ReductionNode.h

171

It's more clear, thanks!

mlir/lib/Reducer/ReductionNode.cpp

64–66

getRange().size() is guaranteed to be >= 1. Yes, it won't cause any problem even it's empty and I agree > 1 will be better.

[Chia-hungDuan]

Addressed review comments

[rriddle]

mlir/include/mlir/Reducer/OptReductionPass.h

0

Do you actually need to declare this? I wouldn't expect it to be necessary.

mlir/lib/Reducer/ReductionTreePass.cpp

110

Why not just use SmallVector instead? It also has the nice pop_back_val API.

162

Do you need the double-ended nature of a queue? Or could you just use a SmallVector based stack?

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

93

What kind of patterns do we expect to be populated here? Would this be better served by a dialect interface of some kind? Where interested dialects could provide patterns for their (or related) operations?

[Chia-hungDuan]

mlir/include/mlir/Reducer/OptReductionPass.h

0

Removed

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

93

In the later CL(D101607) we change mlir-reduce to be a lib, for different project, they can provide their pattern s in the driver, for example, https://reviews.llvm.org/D101607#change-Zg1nOLxd11B7

For example, we may want to reduce a complex tensor into simpler form. I didn't have too many real cases(only got few samples from the codes which stays in the early stage of mlir-reduce), I'm planning to build tf-reduce in tensorflow and see what kinds of patterns may be useful.

[Chia-hungDuan]

Addressed review comment

[rriddle]

mlir/lib/Reducer/Tester.cpp

53

Can you add a comment here?

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

93

I think it would be nice to have a more principled way of providing patterns. The way this is structured right now feels a bit too adhoc. IMO the patterns should either be collected by a dialect interface, users defining their own reduction passes, or some other similar mechanism. Is their something in the structure of mlir-reduce that prevents those types of techniques or makes them harder to integrate with?

[Chia-hungDuan]

Added comments

[jpienaar]

LG, the dialect interface seems like a nice group mechanism

mlir/lib/Reducer/ReductionTreePass.cpp

139

Prefer: /*arg=*/ format to allow clang-tidy warnings if mismatched

162

perhaps marked done accidentally ?

[Chia-hungDuan]

Use Dialect Interface to collect reduce patterns

[Chia-hungDuan]

Verify the module before print it to the file because an invalid module may fail dump.

[rriddle]

Looking good, thanks! May just need one more look over after this round, but should be really close to landing. Awesome work!

mlir/include/mlir/Reducer/OptReductionPass.h

0

Can we delete this file and just move the pass class into the .cpp file?

0

I don't think this include is necessary.

mlir/include/mlir/Reducer/ReducePatternInterface.h

18 ↗

(On Diff #347301)

Can you beef up this documentation? It should be a bit more clear what the contract with the user is, i.e. if I was adding support for my own dialect, how am I supposed to use this interface?

Also, nit: Top-level docs should use ///.

22 ↗

(On Diff #347301)

nit: Can you make this a protected member?

24 ↗

(On Diff #347301)

Can you beef this description up? We should document what types of patterns should be provided here, and give an idea on how they are going to be used.

Also, nit: Top-level docs should use ///.

25 ↗

(On Diff #347301)

nit on name (feel free to ignore): What about populateReductionPatterns?

mlir/include/mlir/Reducer/ReductionNode.h

149

Can you add a doc here?

mlir/include/mlir/Reducer/ReductionTreePass.h

0

Same comment as in OptReductionPass.h: Can we delete this file and move the pattern class into the .cpp file?

mlir/lib/Reducer/ReductionTreePass.cpp

148

Drop the extra //?

149

nit: Can you rename to ReducePatternInterfaceCollection? The current name leads me to believe this is an interface, as opposed to a collection of interfaces. Also, this class should be in an anonymous namespace.

mlir/lib/Reducer/Tester.cpp

31

Do we want to setup a dummy diagnostic handler here? I don't think we want the user to see errors generated here (at least not in the common case).

mlir/test/lib/Dialect/Test/TestInterfaces.h

37 ↗

(On Diff #347301)

nit: Can you move this to TestDialect.cpp? The other test dialect interfaces are there as well: https://github.com/llvm/llvm-project/blob/7a211ed110a72ad453305aba250da50c965c2f8e/mlir/test/lib/Dialect/Test/TestDialect.cpp#L47

mlir/test/lib/Dialect/Test/TestPatterns.cpp

65

Can you just have this be a static method that the interface calls into? That would remove the need to define the interface in the header.

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

16

Is this necessary? Can you trim some of these now?

32

Is this necessary now?

[Chia-hungDuan]

mlir/include/mlir/Reducer/OptReductionPass.h

0

File removed.

0

Done. In fact, I'm thinking to merge this into ReductionTreePass, maybe in the later CL

mlir/include/mlir/Reducer/ReducePatternInterface.h

25 ↗

(On Diff #347301)

Done. In fact, this is an useful comment for me :)

mlir/lib/Reducer/Tester.cpp

31

Yes, I'm going to have another CL for improving the usage, doc also included. This will be covered in that CL

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

32

Sorry, missed to clean this

[Chia-hungDuan]

Addressed review comments and fixed two bugs

1. The reduction path construction should include the root node (Changes in ReductionTreePass.cpp:121)

2. applyOpPatternsAndFold may erase operations, so it couldn't be in the enumerate loop (Changes in ReductionTreePass.cpp:61)

[jpienaar]

Looks good, I'll let River take one more pass

mlir/include/mlir/Reducer/ReductionNode.h

41

s/.etc./etc./

55

Does this effectively mean it returns the most reduced Module that is interesting? E.g., if the pattern hasn't been applied then it is either due to the pattern not being applicable or result not being interesting? Or would it reduce even if the reduced one is no longer interesting? (I don't have much better suggestion, but getModule feels like a const accessor, it doesn't capture much)

78

variants (the child nodes)

91

For passes, we have initialize method, could we follow the same convention/spelling?

mlir/lib/Reducer/ReductionTreePass.cpp

62

Could you document why do we skip checking the result of folding?

72

ReductionNode

[Chia-hungDuan]

mlir/include/mlir/Reducer/ReductionNode.h

55

A ReductionNode will be in either two states (Interestingness::Untested, Interestingness::True/Interestingness::False). If it's tested, the module will have been applied the patterns on the operations within the given ranges and that's the reduced module. The module may not be interesting or has smaller size than others like its parent. It just an exploration state. You're right, I shouldn't say it's reduced, I would say it's applied some reduction strategies.

According to our algorithm, we won't explore an uninteresting state, i.e., ReductionNode with Interestingness::False. If we have a ReductionNode with Interestingness::Untested, which means this node is waiting for testing interestingness.

Added const qualifier and fixed/beefed up the comment.

91

Done.
I used assert for ensuring the result should be success() for now, will update them to proper diagnostic in later CL

[Chia-hungDuan]

Addressed review comments and updated the commit message

[rriddle]

Looks good, thanks!

mlir/include/mlir/Reducer/ReductionPatternInterface.h

44–46

?

47–49

?

mlir/lib/Reducer/OptReductionPass.cpp

80

Only classes should really go in anonymous namespaces, can you move this to the end of the class above?

mlir/lib/Reducer/ReductionTreePass.cpp

35

Only classes should really be in anonymous namespaces, the static functions should be top-scope. Can you restrict the anonymous namespace to just contain classes?

[Chia-hungDuan]

mlir/include/mlir/Reducer/ReductionPatternInterface.h

44–46

Very appreciated for revising the sentences!

[Chia-hungDuan]

Addressed review comments

[Chia-hungDuan]

Shouldn't put statement in assert or it may be removed in certain build.

mlir/include/mlir/Reducer/OptReductionPass.h

-//===- OptReductionPass.h - Optimization Reduction Pass Wrapper -*- 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 Opt Reduction Pass Wrapper. It creates a MLIR pass to
-// run any optimization pass within it and only replaces the output module with
-// the transformed version if it is smaller and interesting.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_REDUCER_OPTREDUCTIONPASS_H
-#define MLIR_REDUCER_OPTREDUCTIONPASS_H
-
-#include "PassDetail.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Pass/PassManager.h"
-#include "mlir/Reducer/ReductionNode.h"
-#include "mlir/Reducer/ReductionTreePass.h"
-#include "mlir/Reducer/Tester.h"
-#include "mlir/Transforms/Passes.h"
-#include "llvm/Support/Debug.h"
-
-namespace mlir {
-
-class OptReductionPass : public OptReductionBase<OptReductionPass> {
-public:
-  OptReductionPass() = default;
-
-  OptReductionPass(const OptReductionPass &srcPass) = default;
-
-  /// Runs the pass instance in the pass pipeline.
-  void runOnOperation() override;
-};
-
-} // end namespace mlir
-
-#endif

mlir/include/mlir/Reducer/Passes.h

 //===- Passes.h - Reducer Pass Construction and Registration ----*- 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
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_REDUCER_PASSES_H
 #define MLIR_REDUCER_PASSES_H
 
 #include "mlir/Pass/Pass.h"
-#include "mlir/Reducer/OptReductionPass.h"
-#include "mlir/Reducer/ReductionTreePass.h"
 
 namespace mlir {
 
 std::unique_ptr<Pass> createReductionTreePass();
 
 std::unique_ptr<Pass> createOptReductionPass();
 
 /// Generate the code for registering reducer passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Reducer/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_REDUCER_PASSES_H

mlir/include/mlir/Reducer/Passes.td

def CommonReductionPassOptions {
   list<Option> options = [
     Option<"testerName", "test", "std::string", /* default */"",
            "The filename of the tester">,
     ListOption<"testerArgs", "test-arg", "std::string",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">,
   ];
 }
 
-def ReductionTree : Pass<"reduction-tree", "ModuleOp"> {
+def ReductionTree : Pass<"reduction-tree"> {
   let summary = "A general reduction tree pass for the MLIR Reduce Tool";
 
   let constructor = "mlir::createReductionTreePass()";
 
   let options = [
-    Option<"opReducerName", "op-reducer", "std::string", /* default */"",
-           "The OpReducer to reduce the module">,
     Option<"traversalModeId", "traversal-mode", "unsigned",
            /* default */"0", "The graph traversal mode">,
   ] # CommonReductionPassOptions.options;
 }
 
 def OptReduction : Pass<"opt-reduction-pass", "ModuleOp"> {
   let summary = "A reduction pass wrapper for optimization passes";
 

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

-//===- OpReducer.h - MLIR Reduce 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.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_REDUCER_PASSES_OPREDUCER_H
-#define MLIR_REDUCER_PASSES_OPREDUCER_H
-
-#include <limits>
-
-#include "mlir/Reducer/ReductionNode.h"
-#include "mlir/Reducer/Tester.h"
-
-namespace mlir {
-
-class OpReducer {
-public:
-  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;
-};
-
-/// Reducer is a helper class to remove potential uninteresting operations from
-/// module.
-template <typename OpType>
-class Reducer : public OpReducer {
-public:
-  ~Reducer() override = default;
-
-  int getNumTargetOps(ModuleOp module) const override {
-    return std::distance(module.getOps<OpType>().begin(),
-                         module.getOps<OpType>().end());
-  }
-
-  void reduce(ModuleOp module,
-              ArrayRef<ReductionNode::Range> rangeToKeep) override {
-    std::vector<Operation *> opsToRemove;
-    size_t keepIndex = 0;
-
-    for (auto op : enumerate(module.getOps<OpType>())) {
-      int index = op.index();
-      if (keepIndex < rangeToKeep.size() &&
-          index == rangeToKeep[keepIndex].second)
-        ++keepIndex;
-      if (keepIndex == rangeToKeep.size() ||
-          index < rangeToKeep[keepIndex].first)
-        opsToRemove.push_back(op.value());
-    }
-
-    for (Operation *o : opsToRemove) {
-      o->dropAllUses();
-      o->erase();
-    }
-  }
-};
-
-} // end namespace mlir
-
-#endif

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 "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 namespace mlir {
 
+class ModuleOp;
+class Region;
+
 /// 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.
+/// ReductionTreePass will build a reduction tree during module reduction and
+/// the ReductionNode represents the vertex of the tree. A ReductionNode records
+/// the information such as the reduced module, how this node is reduced from
+/// the parent node, etc. This information will be used to construct a reduction
+/// path to reduce the certain module.
 class ReductionNode {
 public:
   template <TraversalMode mode>
   class iterator;
 
   using Range = std::pair<int, int>;
 
   ReductionNode(ReductionNode *parent, std::vector<Range> range,
                 llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator);
 
-  ReductionNode *getParent() const;
+  ReductionNode *getParent() const { return parent; }
+
+  /// If the ReductionNode hasn't been tested the interestingness, it'll be the
+  /// same module as the one in the parent node. Otherwise, the returned module
+  /// will have been applied certain reduction strategies. Note that it's not
+  /// necessary to be an interesting case or a reduced module (has smaller size
+  /// than parent's).
+  ModuleOp getModule() const { return module; }
+
+  /// Return the region we're reducing.
+  Region &getRegion() const { return *region; }
 
-  size_t getSize() const;
+  /// Return the size of the module.
+  size_t getSize() const { return size; }
 
   /// Returns true if the module exhibits the interesting behavior.
-  Tester::Interestingness isInteresting() const;
+  Tester::Interestingness isInteresting() const { return interesting; }
 
-  std::vector<Range> getRanges() const;
+  /// Return the range information that how this node is reduced from the parent
+  /// node.
+  ArrayRef<Range> getStartRanges() const { return startRanges; }
 
-  std::vector<ReductionNode *> &getVariants();
+  /// Return the range set we are using to generate variants.
+  ArrayRef<Range> getRanges() const { return ranges; }
+
+  /// Return the generated variants(the child nodes).
+  ArrayRef<ReductionNode *> getVariants() const { return variants; }
 
   /// Split the ranges and generate new variants.
-  std::vector<ReductionNode *> generateNewVariants();
+  ArrayRef<ReductionNode *> generateNewVariants();
 
   /// Update the interestingness result from tester.
   void update(std::pair<Tester::Interestingness, size_t> result);
 
+  /// Each Reduction Node contains a copy of module for applying rewrite
+  /// patterns. In addition, we only apply rewrite patterns in a certain region.
+  /// In init(), we will duplicate the module from parent node and locate the
+  /// corresponding region.
+  LogicalResult initialize(ModuleOp parentModule, Region &parentRegion);
+
 private:
   /// 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)
+      for (ReductionNode *node : getNeighbors(top))
         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) {
+    ArrayRef<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.
+  /// This is a copy of module from parent node. All the reducer patterns will
+  /// be applied to this instance.
+  ModuleOp module;
+
+  /// The region of certain operation we're reducing in the module
+  Region *region;
+
+  /// The node we are reduced from. It means we will be in variants of parent
+  /// node.
+  ReductionNode *parent;
+
+  /// The size of module after applying the reducer patterns with range
+  /// constraints. This is only valid while the interestingness has been tested.
   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.
+  /// `ranges` represents the selected subset of operations in the region. We
+  /// implictly number each operation in the region and ReductionTreePass will
+  /// apply reducer patterns on the operation falls into the `ranges`. We will
+  /// generate new ReductionNode with subset of `ranges` to see if we can do
+  /// further reduction. we may split the element in the `ranges` so that we can
+  /// have more subset variants from `ranges`.
+  /// Note that after applying the reducer patterns the number of operation in
+  /// the region may have changed, we need to update the `ranges` after that.
   std::vector<Range> ranges;
 
+  /// `startRanges` records the ranges of operations selected from the parent
+  /// node to produce this ReductionNode. It can be used to construct the
+  /// reduction path from the root. I.e., if we apply the same reducer patterns
+  /// and `startRanges` selection on the parent region, we will get the same
+  /// module as this node.
+  const std::vector<Range> startRanges;
+
   /// This points to the child variants that were created using this node as a
   /// starting point.
   std::vector<ReductionNode *> variants;
 
   llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator;
 };
 
 // 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);
+  ArrayRef<ReductionNode *> getNeighbors(ReductionNode *node);
 };
 
 } // end namespace mlir
 
-#endif
+#endif // MLIR_REDUCER_REDUCTIONNODE_H

mlir/include/mlir/Reducer/ReductionPatternInterface.h

+//===- ReducePatternInterface.h - Collecting Reduce Patterns ----*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_REDUCER_REDUCTIONPATTERNINTERFACE_H
+#define MLIR_REDUCER_REDUCTIONPATTERNINTERFACE_H
+
+#include "mlir/IR/DialectInterface.h"
+
+namespace mlir {
+
+class RewritePatternSet;
+
+/// This is used to report the reduction patterns for a Dialect. While using
+/// mlir-reduce to reduce a module, we may want to transform certain cases into
+/// simpler forms by applying certain rewrite patterns. Implement the
+/// `populateReductionPatterns` to report those patterns by adding them to the
+/// RewritePatternSet.
+///
+/// Example:
+///   MyDialectReductionPattern::populateReductionPatterns(
+///       RewritePatternSet &patterns) {
+///       patterns.add<TensorOpReduction>(patterns.getContext());
+///   }
+///
+/// For DRR, mlir-tblgen will generate a helper function
+/// `populateWithGenerated` which has the same signature therefore you can
+/// delegate to the helper function as well.
+///
+/// Example:
+///   MyDialectReductionPattern::populateReductionPatterns(
+///       RewritePatternSet &patterns) {
+///       // Include the autogen file somewhere above.
+///       populateWithGenerated(patterns);
+///   }
+class DialectReductionPatternInterface
+    : public DialectInterface::Base<DialectReductionPatternInterface> {
+public:
+  /// Patterns provided here are intended to transform operations from a complex
+  /// form to a simpler form, without breaking the semantics of the program
+  /// being reduced. For example, you may want to replace the
+  /// tensor<?xindex> with a known rank and type, e.g. tensor<1xi32>, or
+  /// replacing an operation with a constant.
+  virtual void populateReductionPatterns(RewritePatternSet &patterns) const = 0;
+
+protected:
+  DialectReductionPatternInterface(Dialect *dialect) : Base(dialect) {}
+};
+
+} // end namespace mlir
+
+#endif // MLIR_REDUCER_REDUCTIONPATTERNINTERFACE_H

mlir/include/mlir/Reducer/ReductionTreePass.h

-//===- ReductionTreePass.h - Reduction Tree Pass Implementation -*- 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 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.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_REDUCER_REDUCTIONTREEPASS_H
-#define MLIR_REDUCER_REDUCTIONTREEPASS_H
-
-#include <vector>
-
-#include "PassDetail.h"
-#include "ReductionNode.h"
-#include "mlir/Reducer/Passes/OpReducer.h"
-#include "mlir/Reducer/Tester.h"
-
-#define DEBUG_TYPE "mlir-reduce"
-
-namespace mlir {
-
-/// 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.
-class ReductionTreePass : public ReductionTreeBase<ReductionTreePass> {
-public:
-  ReductionTreePass() = default;
-  ReductionTreePass(const ReductionTreePass &pass) = default;
-
-  /// Runs the pass instance in the pass pipeline.
-  void runOnOperation() override;
-
-private:
-  template <typename IteratorType>
-  ModuleOp findOptimal(ModuleOp module, std::unique_ptr<OpReducer> reducer,
-                       ReductionNode *node);
-};
-
-} // end namespace mlir
-
-#endif

mlir/lib/Reducer/CMakeLists.txt

 add_mlir_library(MLIRReduce
+   OptReductionPass.cpp
+   ReductionNode.cpp
+   ReductionTreePass.cpp
    Tester.cpp
    LINK_LIBS PUBLIC
    MLIRIR
+   MLIRPass
+   MLIRRewrite
+   MLIRTransformUtils
 )
 
 mlir_check_all_link_libraries(MLIRReduce)

mlir/lib/Reducer/OptReductionPass.cpp

 //===- OptReductionPass.cpp - Optimization Reduction Pass Wrapper ---------===//
 //
 // 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 Opt Reduction Pass Wrapper. It creates a MLIR pass to
 // run any optimization pass within it and only replaces the output module with
 // the transformed version if it is smaller and interesting.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Reducer/OptReductionPass.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassRegistry.h"
+#include "mlir/Reducer/PassDetail.h"
 #include "mlir/Reducer/Passes.h"
 #include "mlir/Reducer/Tester.h"
+#include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "mlir-reduce"
 
 using namespace mlir;
 
+namespace {
+
+class OptReductionPass : public OptReductionBase<OptReductionPass> {
+public:
+  /// Runs the pass instance in the pass pipeline.
+  void runOnOperation() override;
+};
+
+} // end anonymous namespace
+
 /// Runs the pass instance in the pass pipeline.
 void OptReductionPass::runOnOperation() {
   LLVM_DEBUG(llvm::dbgs() << "\nOptimization Reduction pass: ");
 
   Tester test(testerName, testerArgs);
 
   ModuleOp module = this->getOperation();
   ModuleOp moduleVariant = module.clone();
if (reduced.first == Tester::Interestingness::True &&
     LLVM_DEBUG(llvm::dbgs() << "\nUnsuccessful Transformed version\n\n");
   }
 
   moduleVariant->destroy();
 
   LLVM_DEBUG(llvm::dbgs() << "Pass Complete\n\n");
 }
 
 std::unique_ptr<Pass> mlir::createOptReductionPass() {
   return std::make_unique<OptReductionPass>();
 }

mlir/lib/Reducer/ReductionNode.cpp

 // metadata for a specific generated variant within a reduction pass and are the
 // building blocks of the reduction tree structure. A reduction tree is used to
 // keep track of the different generated variants throughout a reduction pass in
 // the MLIR Reduce tool.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Reducer/ReductionNode.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "llvm/ADT/STLExtras.h"
 
 #include <algorithm>
 #include <limits>
 
 using namespace mlir;
 
 ReductionNode::ReductionNode(
-    ReductionNode *parent, std::vector<Range> ranges,
+    ReductionNode *parentNode, std::vector<Range> ranges,
     llvm::SpecificBumpPtrAllocator<ReductionNode> &allocator)
-    : size(std::numeric_limits<size_t>::max()),
-      interesting(Tester::Interestingness::Untested),
-      /// Root node will have the parent pointer point to themselves.
-      parent(parent == nullptr ? this : parent), ranges(ranges),
-      allocator(allocator) {}
-
-/// Returns the size in bytes of the module.
-size_t ReductionNode::getSize() const { return size; }
-
-ReductionNode *ReductionNode::getParent() const { return parent; }
-
-/// Returns true if the module exhibits the interesting behavior.
-Tester::Interestingness ReductionNode::isInteresting() const {
-  return interesting;
+    : parent(parentNode == nullptr ? this : parentNode),
+      size(std::numeric_limits<size_t>::max()),
+      interesting(Tester::Interestingness::Untested), ranges(ranges),
+      startRanges(ranges), allocator(allocator) {
+  if (parent != this)
+    if (failed(initialize(parent->getModule(), parent->getRegion())))
+      llvm_unreachable("unexpected initialization failure");
+}
+
+LogicalResult ReductionNode::initialize(ModuleOp parentModule,
+                                        Region &targetRegion) {
+  // Use the mapper help us find the corresponding region after module clone.
+  BlockAndValueMapping mapper;
+  module = cast<ModuleOp>(parentModule->clone(mapper));
+  // Use the first block of targetRegion to locate the cloned region.
+  Block *block = mapper.lookup(&*targetRegion.begin());
+  region = block->getParent();
+  return success();
 }
 
-std::vector<ReductionNode::Range> ReductionNode::getRanges() const {
-  return ranges;
-}
-
-std::vector<ReductionNode *> &ReductionNode::getVariants() { return variants; }
-
-#include <iostream>
-
 /// If we haven't explored any variants from this node, we will create N
 /// variants, N is the length of `ranges` if N > 1. Otherwise, we will split the
 /// max element in `ranges` and create 2 new variants for each call.
-std::vector<ReductionNode *> ReductionNode::generateNewVariants() {
-  std::vector<ReductionNode *> newNodes;
+ArrayRef<ReductionNode *> ReductionNode::generateNewVariants() {
+  int oldNumVariant = getVariants().size();
+
+  auto createNewNode = [this](std::vector<Range> ranges) {
+    return new (allocator.Allocate())
+        ReductionNode(this, std::move(ranges), allocator);
+  };
 
   // If we haven't created new variant, then we can create varients by removing
   // each of them respectively. For example, given {{1, 3}, {4, 9}}, we can
   // produce variants with range {{1, 3}} and {{4, 9}}.
-  if (variants.size() == 0 && ranges.size() != 1) {
-    for (const Range &range : ranges) {
-      std::vector<Range> subRanges = ranges;
+  if (variants.size() == 0 && getRanges().size() > 1) {
+    for (const Range &range : getRanges()) {
+      std::vector<Range> subRanges = getRanges();
       llvm::erase_value(subRanges, range);
-      ReductionNode *newNode = allocator.Allocate();
-      new (newNode) ReductionNode(this, subRanges, allocator);
-      newNodes.push_back(newNode);
-      variants.push_back(newNode);
+      variants.push_back(createNewNode(std::move(subRanges)));
     }
 
-    return newNodes;
+    return getVariants().drop_front(oldNumVariant);
   }
 
   // At here, we have created the type of variants mentioned above. We would
   // like to split the max range into 2 to create 2 new variants. Continue on
   // the above example, we split the range {4, 9} into {4, 6}, {6, 9}, and
   // create two variants with range {{1, 3}, {4, 6}} and {{1, 3}, {6, 9}}. The
-  // result ranges vector will be {{1, 3}, {4, 6}, {6, 9}}.
+  // final ranges vector will be {{1, 3}, {4, 6}, {6, 9}}.
   auto maxElement = std::max_element(
       ranges.begin(), ranges.end(), [](const Range &lhs, const Range &rhs) {
         return (lhs.second - lhs.first) > (rhs.second - rhs.first);
       });
 
-  // We can't split range with lenght 1, which means we can't produce new
+  // The length of range is less than 1, we can't split it to create new
   // variant.
-  if (maxElement->second - maxElement->first == 1)
+  if (maxElement->second - maxElement->first <= 1)
     return {};
 
-  auto createNewNode = [this](const std::vector<Range> &ranges) {
-    ReductionNode *newNode = allocator.Allocate();
-    new (newNode) ReductionNode(this, ranges, allocator);
-    return newNode;
-  };
-
   Range maxRange = *maxElement;
-  std::vector<Range> subRanges = ranges;
+  std::vector<Range> subRanges = getRanges();
   auto subRangesIter = subRanges.begin() + (maxElement - ranges.begin());
   int half = (maxRange.first + maxRange.second) / 2;
   *subRangesIter = std::make_pair(maxRange.first, half);
-  newNodes.push_back(createNewNode(subRanges));
+  variants.push_back(createNewNode(subRanges));
   *subRangesIter = std::make_pair(half, maxRange.second);
-  newNodes.push_back(createNewNode(subRanges));
+  variants.push_back(createNewNode(std::move(subRanges)));
 
-  variants.insert(variants.end(), newNodes.begin(), newNodes.end());
   auto it = ranges.insert(maxElement, std::make_pair(half, maxRange.second));
   it = ranges.insert(it, std::make_pair(maxRange.first, half));
   // Remove the range that has been split.
   ranges.erase(it + 2);
 
-  return newNodes;
+  return getVariants().drop_front(oldNumVariant);
 }
 
 void ReductionNode::update(std::pair<Tester::Interestingness, size_t> result) {
   std::tie(interesting, size) = result;
+  // After applying reduction, the number of operation in the region may have
+  // changed. Non-interesting case won't be explored thus it's safe to keep it
+  // in a stale status.
+  if (interesting == Tester::Interestingness::True) {
+    // This module may has been updated. Reset the range.
+    ranges.clear();
+    ranges.push_back({0, std::distance(region->op_begin(), region->op_end())});
+  }
 }
 
-std::vector<ReductionNode *>
+ArrayRef<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.
-  llvm::ArrayRef<ReductionNode *> variantsFromParent =
+  ArrayRef<ReductionNode *> variantsFromParent =
       node->getParent()->getVariants();
 
   // The parent node created several variants and they may be waiting for
   // examing interestingness. In Single Path approach, we will select the
   // smallest variant to continue our exploration. Thus we should wait until the
   // last variant to be examed then do the following traversal decision.
   if (!llvm::all_of(variantsFromParent, [](ReductionNode *node) {
         return node->isInteresting() != Tester::Interestingness::Untested;
       })) {
     return {};
   }
 
   ReductionNode *smallest = nullptr;
   for (ReductionNode *node : variantsFromParent) {
     if (node->isInteresting() != Tester::Interestingness::True)
       continue;
     if (smallest == nullptr || node->getSize() < smallest->getSize())
       smallest = node;
   }
 
-  if (smallest != nullptr) {
+  if (smallest != nullptr &&
+      smallest->getSize() < node->getParent()->getSize()) {
     // 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();
   }
 
   return node->generateNewVariants();
 }

mlir/lib/Reducer/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/IR/DialectInterface.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Reducer/PassDetail.h"
+#include "mlir/Reducer/Passes.h"
+#include "mlir/Reducer/ReductionNode.h"
+#include "mlir/Reducer/ReductionPatternInterface.h"
+#include "mlir/Reducer/Tester.h"
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/ManagedStatic.h"
+
+using namespace mlir;
+
+/// We implicitly number each operation in the region and if an operation's
+/// number falls into rangeToKeep, we need to keep it and apply the given
+/// rewrite patterns on it.
+static void applyPatterns(Region &region,
+                          const FrozenRewritePatternSet &patterns,
+                          ArrayRef<ReductionNode::Range> rangeToKeep,
+                          bool eraseOpNotInRange) {
+  std::vector<Operation *> opsNotInRange;
+  std::vector<Operation *> opsInRange;
+  size_t keepIndex = 0;
+  for (auto op : enumerate(region.getOps())) {
+    int index = op.index();
+    if (keepIndex < rangeToKeep.size() &&
+        index == rangeToKeep[keepIndex].second)
+      ++keepIndex;
+    if (keepIndex == rangeToKeep.size() || index < rangeToKeep[keepIndex].first)
+      opsNotInRange.push_back(&op.value());
+    else
+      opsInRange.push_back(&op.value());
+  }
+
+  // `applyOpPatternsAndFold` may erase the ops so we can't do the pattern
+  // matching in above iteration. Besides, erase op not-in-range may end up in
+  // invalid module, so `applyOpPatternsAndFold` should come before that
+  // transform.
+  for (Operation *op : opsInRange)
+    // `applyOpPatternsAndFold` returns whether the op is convered. Omit it
+    // because we don't have expectation this reduction will be success or not.
+    (void)applyOpPatternsAndFold(op, patterns);
+
+  if (eraseOpNotInRange)
+    for (Operation *op : opsNotInRange) {
+      op->dropAllUses();
+      op->erase();
+    }
+}
+
+/// We will apply the reducer patterns to the operations in the ranges specified
+/// by ReductionNode. Note that we are not able to remove an operation without
+/// replacing it with another valid operation. However, The validity of module
+/// reduction is based on the Tester provided by the user and that means certain
+/// invalid module is still interested by the use. Thus we provide an
+/// alternative way to remove operations, which is using `eraseOpNotInRange` to
+/// erase the operations not in the range specified by ReductionNode.
+template <typename IteratorType>
+static void findOptimal(ModuleOp module, Region &region,
+                        const FrozenRewritePatternSet &patterns,
+                        const Tester &test, bool eraseOpNotInRange) {
+  std::pair<Tester::Interestingness, size_t> initStatus =
+      test.isInteresting(module);
+  // While exploring the reduction tree, we always branch from an interesting
+  // node. Thus the root node must be interesting.
+  if (initStatus.first != Tester::Interestingness::True)
+    return;
+
+  llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;
+
+  std::vector<ReductionNode::Range> ranges{
+      {0, std::distance(region.op_begin(), region.op_end())}};
+
+  ReductionNode *root = allocator.Allocate();
+  new (root) ReductionNode(nullptr, std::move(ranges), allocator);
+  // Duplicate the module for root node and locate the region in the copy.
+  if (failed(root->initialize(module, region)))
+    llvm_unreachable("unexpected initialization failure");
+  root->update(initStatus);
+
+  ReductionNode *smallestNode = root;
+  IteratorType iter(root);
+
+  while (iter != IteratorType::end()) {
+    ReductionNode &currentNode = *iter;
+    Region &curRegion = currentNode.getRegion();
+
+    applyPatterns(curRegion, patterns, currentNode.getRanges(),
+                  eraseOpNotInRange);
+    currentNode.update(test.isInteresting(currentNode.getModule()));
+
+    if (currentNode.isInteresting() == Tester::Interestingness::True &&
+        currentNode.getSize() < smallestNode->getSize())
+      smallestNode = &currentNode;
+
+    ++iter;
+  }
+
+  // At here, we have found an optimal path to reduce the given region. Retrieve
+  // the path and apply the reducer to it.
+  SmallVector<ReductionNode *> trace;
+  ReductionNode *curNode = smallestNode;
+  trace.push_back(curNode);
+  while (curNode != root) {
+    curNode = curNode->getParent();
+    trace.push_back(curNode);
+  }
+
+  // Reduce the region through the optimal path.
+  while (!trace.empty()) {
+    ReductionNode *top = trace.pop_back_val();
+    applyPatterns(region, patterns, top->getStartRanges(), eraseOpNotInRange);
+  }
+
+  if (test.isInteresting(module).first != Tester::Interestingness::True)
+    llvm::report_fatal_error("Reduced module is not interesting");
+  if (test.isInteresting(module).second != smallestNode->getSize())
+    llvm::report_fatal_error(
+        "Reduced module doesn't have consistent size with smallestNode");
+}
+
+template <typename IteratorType>
+static void findOptimal(ModuleOp module, Region &region,
+                        const FrozenRewritePatternSet &patterns,
+                        const Tester &test) {
+  // We separate the reduction process into 2 steps, the first one is to erase
+  // redundant operations and the second one is to apply the reducer patterns.
+
+  // In the first phase, we don't apply any patterns so that we only select the
+  // range of operations to keep to the module stay interesting.
+  findOptimal<IteratorType>(module, region, /*patterns=*/{}, test,
+                            /*eraseOpNotInRange=*/true);
+  // In the second phase, we suppose that no operation is redundant, so we try
+  // to rewrite the operation into simpler form.
+  findOptimal<IteratorType>(module, region, patterns, test,
+                            /*eraseOpNotInRange=*/false);
+}
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Reduction Pattern Interface Collection
+//===----------------------------------------------------------------------===//
+
+class ReductionPatternInterfaceCollection
+    : public DialectInterfaceCollection<DialectReductionPatternInterface> {
+public:
+  using Base::Base;
+
+  // Collect the reduce patterns defined by each dialect.
+  void populateReductionPatterns(RewritePatternSet &pattern) const {
+    for (const DialectReductionPatternInterface &interface : *this)
+      interface.populateReductionPatterns(pattern);
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// ReductionTreePass
+//===----------------------------------------------------------------------===//
+
+/// 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.
+class ReductionTreePass : public ReductionTreeBase<ReductionTreePass> {
+public:
+  ReductionTreePass() = default;
+  ReductionTreePass(const ReductionTreePass &pass) = default;
+
+  LogicalResult initialize(MLIRContext *context) override;
+
+  /// Runs the pass instance in the pass pipeline.
+  void runOnOperation() override;
+
+private:
+  void reduceOp(ModuleOp module, Region &region);
+
+  FrozenRewritePatternSet reducerPatterns;
+};
+
+} // end anonymous namespace
+
+LogicalResult ReductionTreePass::initialize(MLIRContext *context) {
+  RewritePatternSet patterns(context);
+  ReductionPatternInterfaceCollection reducePatternCollection(context);
+  reducePatternCollection.populateReductionPatterns(patterns);
+  reducerPatterns = std::move(patterns);
+  return success();
+}
+
+void ReductionTreePass::runOnOperation() {
+  Operation *topOperation = getOperation();
+  while (topOperation->getParentOp() != nullptr)
+    topOperation = topOperation->getParentOp();
+  ModuleOp module = cast<ModuleOp>(topOperation);
+
+  SmallVector<Operation *, 8> workList;
+  workList.push_back(getOperation());
+
+  do {
+    Operation *op = workList.pop_back_val();
+
+    for (Region &region : op->getRegions())
+      if (!region.empty())
+        reduceOp(module, region);
+
+    for (Region &region : op->getRegions())
+      for (Operation &op : region.getOps())
+        if (op.getNumRegions() != 0)
+          workList.push_back(&op);
+  } while (!workList.empty());
+}
+
+void ReductionTreePass::reduceOp(ModuleOp module, Region &region) {
+  Tester test(testerName, testerArgs);
+  switch (traversalModeId) {
+  case TraversalMode::SinglePath:
+    findOptimal<ReductionNode::iterator<TraversalMode::SinglePath>>(
+        module, region, reducerPatterns, test);
+    break;
+  default:
+    llvm_unreachable("Unsupported mode");
+  }
+}
+
+std::unique_ptr<Pass> mlir::createReductionTreePass() {
+  return std::make_unique<ReductionTreePass>();
+}

mlir/lib/Reducer/Tester.cpp

 //
 // A Tester object is passed as an argument to the reduction passes and it is
 // used to run the interestingness testing script on the different generated
 // reduced variants of the test case.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Reducer/Tester.h"
-
+#include "mlir/IR/Verifier.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, size_t>
 Tester::isInteresting(ModuleOp module) const {
+  // The reduced module should always be vaild, or we may end up retaining the
+  // error message by an invalid case. Besides, an invalid module may not be
+  // able to print properly.
+  if (failed(verify(module)))
+    return std::make_pair(Interestingness::False, /*size=*/0);
+
   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 the IR to file '" + filepath);
 
   size_t 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.
 Tester::Interestingness Tester::isInteresting(StringRef testCase) const {
-
   std::vector<StringRef> testerArgs;
   testerArgs.push_back(testCase);
 
   for (const std::string &arg : testScriptArgs)
     testerArgs.push_back(arg);
 
   testerArgs.push_back(testCase);
 

mlir/test/lib/Dialect/Test/CMakeLists.txt

add_mlir_library(MLIRTestDialect
   MLIRDataLayoutInterfaces
   MLIRDerivedAttributeOpInterface
   MLIRDialect
   MLIRDLTI
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRLinalgTransforms
   MLIRPass
+  MLIRReduce
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRTransformUtils
   MLIRTransforms
 )

mlir/test/lib/Dialect/Test/TestDialect.h

 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 class DLTIDialect;
+class RewritePatternSet;
 } // namespace mlir
 
 #include "TestOpEnums.h.inc"
 #include "TestOpInterfaces.h.inc"
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace mlir {
 namespace test {
 void registerTestDialect(DialectRegistry &registry);
+void populateTestReductionPatterns(RewritePatternSet &patterns);
 } // namespace test
 } // namespace mlir
 
 #endif // MLIR_TESTDIALECT_H

mlir/test/lib/Dialect/Test/TestDialect.cpp

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "TestAttributes.h"
+#include "TestInterfaces.h"
 #include "TestTypes.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 using namespace mlir::test;
 
 void mlir::test::registerTestDialect(DialectRegistry &registry) {
Operation *materializeCallConversion(OpBuilder &builder, Value input,
     if (!(resultType.isSignlessInteger(16) ||
           resultType.isSignlessInteger(32)) ||
         !(input.getType().isSignlessInteger(16) ||
           input.getType().isSignlessInteger(32)))
       return nullptr;
     return builder.create<TestCastOp>(conversionLoc, resultType, input);
   }
 };
+
+struct TestReductionPatternInterface : public DialectReductionPatternInterface {
+public:
+  TestReductionPatternInterface(Dialect *dialect)
+      : DialectReductionPatternInterface(dialect) {}
+
+  virtual void
+  populateReductionPatterns(RewritePatternSet &patterns) const final {
+    populateTestReductionPatterns(patterns);
+  }
+};
+
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 static void testSideEffectOpGetEffect(
     Operation *op,

 void TestDialect::initialize() {
   registerAttributes();
   registerTypes();
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
-                TestInlinerInterface>();
+                TestInlinerInterface, TestReductionPatternInterface>();
   allowUnknownOperations();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.
   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
 }
 TestDialect::~TestDialect() {
   delete static_cast<TestOpEffectInterfaceFallback *>(

mlir/test/lib/Dialect/Test/TestOps.td

let description = [{
     The data layout query pass pattern-matches this op and attaches to it an
     array attribute containing the result of data layout query of the result
     type of this op.
   }];
 
   let results = (outs AnyType:$res);
 }
 
+//===----------------------------------------------------------------------===//
+// Test Reducer Patterns
+//===----------------------------------------------------------------------===//
+
+def OpCrashLong : TEST_Op<"op_crash_long"> {
+  let arguments = (ins I32, I32, I32);
+  let results = (outs I32);
+}
+
+def OpCrashShort : TEST_Op<"op_crash_short"> {
+  let results = (outs I32);
+}
+
+def : Pat<(OpCrashLong $_, $_, $_), (OpCrashShort)>;
+
 #endif // TEST_OPS

mlir/test/lib/Dialect/Test/TestPatterns.cpp

static Attribute OpMTest(PatternRewriter &rewriter, Value val) {
   return rewriter.getIntegerAttr(rewriter.getIntegerType(32), i);
 }
 
 namespace {
 #include "TestPatterns.inc"
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
+// Test Reduce Pattern Interface
+//===----------------------------------------------------------------------===//
+
+void mlir::test::populateTestReductionPatterns(RewritePatternSet &patterns) {
+  populateWithGenerated(patterns);
+}
+
+//===----------------------------------------------------------------------===//
 // Canonicalizer Driver.
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct FoldingPattern : public RewritePattern {
 public:
   FoldingPattern(MLIRContext *context)
       : RewritePattern(TestOpInPlaceFoldAnchor::getOperationName(),

mlir/test/lib/Reducer/MLIRTestReducer.cpp

 
 } // end anonymous namespace
 
 void TestReducer::runOnFunction() {
   for (auto &op : getOperation())
     op.walk([&](Operation *op) {
       StringRef opName = op->getName().getStringRef();
 
-      if (opName == "test.crashOp") {
+      if (opName.contains("op_crash")) {
         llvm::errs() << "MLIR Reducer Test generated failure: Found "
                         "\"crashOp\" operation\n";
         exit(1);
       }
     });
 }
 
 namespace mlir {
 void registerTestReducer() {
   PassRegistration<TestReducer>(
       PASS_NAME, "Tests MLIR Reduce tool by generating failures");
 }
 } // namespace mlir

mlir/test/mlir-reduce/crashop-reduction.mlir

+// UNSUPPORTED: system-windows
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/failure-test.sh' | FileCheck %s
+// "test.op_crash_long" should be replaced with a shorter form "test.op_crash_short".
+
+// CHECK-NOT: func @simple1() {
+func @simple1() {
+  return
+}
+
+// CHECK-LABEL: func @simple2(%arg0: i32, %arg1: i32, %arg2: i32) {
+func @simple2(%arg0: i32, %arg1: i32, %arg2: i32) {
+  // CHECK-LABEL: %0 = "test.op_crash_short"() : () -> i32
+  %0 = "test.op_crash_long" (%arg0, %arg1, %arg2) : (i32, i32, i32) -> i32
+  return
+}
+
+// CHECK-NOT: func @simple5() {
+func @simple5() {
+  return
+}

mlir/test/mlir-reduce/dce-test.mlir

 // UNSUPPORTED: system-windows
 // RUN: mlir-reduce %s -opt-reduction-pass='opt-pass=symbol-dce test=%S/failure-test.sh' | FileCheck %s
 // This input should be reduced by the pass pipeline so that only
 // the @simple1 function remains as the other functions should be
 // removed by the dead code elimination pass.
 
 // CHECK-NOT: func private @dead_private_function
 func private @dead_private_function()
 
 // CHECK-NOT: func nested @dead_nested_function
 func nested @dead_nested_function()
 
 // CHECK-LABEL: func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
-  "test.crashOp" () : () -> ()
+  "test.op_crash" () : () -> ()
   return
 }

mlir/test/mlir-reduce/multiple-function.mlir

 // UNSUPPORTED: system-windows
-// RUN: mlir-reduce %s -reduction-tree='op-reducer=func traversal-mode=0 test=%S/failure-test.sh' | FileCheck %s
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/failure-test.sh' | FileCheck %s
 // This input should be reduced by the pass pipeline so that only
 // the @simple5 function remains as this is the shortest function
 // containing the interesting behavior.
 
 // CHECK-NOT: func @simple1() {
 func @simple1() {
   return
 }
 
 // CHECK-NOT: func @simple2() {
 func @simple2() {
   return
 }
 
 // CHECK-LABEL: func @simple3() {
 func @simple3() {
-  "test.crashOp" () : () -> ()
+  "test.op_crash" () : () -> ()
   return
 }
 
 // CHECK-NOT: func @simple4() {
 func @simple4(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
-  "test.crashOp"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+  "test.op_crash"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
 // CHECK-NOT: func @simple5() {
 func @simple5() {
   return
 }

mlir/test/mlir-reduce/simple-test.mlir

 // UNSUPPORTED: system-windows
-// RUN: mlir-reduce %s -reduction-tree='op-reducer=func traversal-mode=0 test=%S/test.sh'
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/test.sh'
 
 func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   return
 }

mlir/test/mlir-reduce/single-function.mlir

 // RUN: mlir-opt %s
 // RUN: not mlir-opt %s -test-mlir-reducer -pass-test function-reducer
 
 func @test() {
-  "test.crashOp"() : () -> ()
+  "test.op_crash"() : () -> ()
   return
 }

mlir/tools/mlir-reduce/CMakeLists.txt

set(LIBS
   MLIRPass
   MLIRReduce
   MLIRSupport
   MLIRTransforms
   MLIRTransformUtils
   )
 
 add_llvm_tool(mlir-reduce
-  OptReductionPass.cpp
-  ReductionNode.cpp
-  ReductionTreePass.cpp
   mlir-reduce.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Reducer
 
   DEPENDS
   MLIRReducerIncGen
   )
 
 target_link_libraries(mlir-reduce PRIVATE ${LIBS})
 llvm_update_compile_flags(mlir-reduce)
 
 mlir_check_all_link_libraries(mlir-reduce)

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 "mlir/Reducer/Passes.h"
-
-#include "llvm/Support/Allocator.h"
-
-using namespace mlir;
-
-static 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");
-}
-
-void ReductionTreePass::runOnOperation() {
-  ModuleOp module = this->getOperation();
-  std::unique_ptr<OpReducer> reducer = getOpReducer(opReducerName);
-  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 (traversalModeId) {
-  case TraversalMode::SinglePath:
-    golden = findOptimal<ReductionNode::iterator<TraversalMode::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 <typename IteratorType>
-ModuleOp ReductionTreePass::findOptimal(ModuleOp module,
-                                        std::unique_ptr<OpReducer> reducer,
-                                        ReductionNode *root) {
-  Tester test(testerName, testerArgs);
-  std::pair<Tester::Interestingness, size_t> initStatus =
-      test.isInteresting(module);
-
-  if (initStatus.first != Tester::Interestingness::True) {
-    LLVM_DEBUG(llvm::dbgs() << "\nThe original input is not interested");
-    return module;
-  }
-
-  root->update(initStatus);
-
-  ReductionNode *smallestNode = root;
-  ModuleOp golden = module;
-
-  IteratorType iter(root);
-
-  while (iter != IteratorType::end()) {
-    ModuleOp cloneModule = module.clone();
-
-    ReductionNode &currentNode = *iter;
-    reducer->reduce(cloneModule, currentNode.getRanges());
-
-    std::pair<Tester::Interestingness, size_t> 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;
-}
-
-std::unique_ptr<Pass> mlir::createReductionTreePass() {
-  return std::make_unique<ReductionTreePass>();
-}

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

 //===- mlir-reduce.cpp - The MLIR reducer ---------------------------------===//
 //
 // 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 implements the general framework of the MLIR reducer tool. It
 // parses the command line arguments, parses the initial MLIR test case and sets
 // up the testing environment. It  outputs the most reduced test case variant
 // after executing the reduction passes.
 //
 //===----------------------------------------------------------------------===//
 
-#include <vector>
-
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
-#include "mlir/Reducer/OptReductionPass.h"
 #include "mlir/Reducer/Passes.h"
-#include "mlir/Reducer/Passes/OpReducer.h"
-#include "mlir/Reducer/ReductionNode.h"
-#include "mlir/Reducer/ReductionTreePass.h"
-#include "mlir/Reducer/Tester.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
-#include "mlir/Transforms/Passes.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
 namespace mlir {
 namespace test {
 void registerTestDialect(DialectRegistry &);
 } // namespace test
 } // namespace mlir
 
 static llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional,
                                                 llvm::cl::Required,
                                                 llvm::cl::desc("<input file>"));
 
 static llvm::cl::opt<std::string>
     outputFilename("o",
if (failed(loadModule(context, moduleRef, inputFilename)))
     llvm::report_fatal_error("Input test case can't be parsed");
 
   auto errorHandler = [&](const Twine &msg) {
     return emitError(UnknownLoc::get(&context)) << msg;
   };
 
   // Reduction pass pipeline.
   PassManager pm(&context);
   if (failed(parser.addToPipeline(pm, errorHandler)))
     llvm::report_fatal_error("Failed to add pipeline");
 
   ModuleOp m = moduleRef.get().clone();
 
   if (failed(pm.run(m)))
     llvm::report_fatal_error("Error running the reduction pass pipeline");
 
   m.print(output->os());
   output->keep();
 
   return 0;
 }