[mlir-reduce] Fix the memory leak and recycle unused modules.

mlir/include/mlir/Reducer/ReductionNode.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_REDUCER_REDUCTIONNODE_H
 #define MLIR_REDUCER_REDUCTIONNODE_H
 
 #include <queue>
 #include <vector>
 
+#include "mlir/IR/OwningOpRef.h"
 #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 {
 
public:
 
   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; }
+  ModuleOp getModule() const { return module.get(); }
 
   /// Return the region we're reducing.
   Region &getRegion() const { return *region; }
 
   /// Return the size of the module.
   size_t getSize() const { return size; }
 
   /// Returns true if the module exhibits the interesting behavior.
protected:
     }
 
   private:
     std::queue<ReductionNode *> visitQueue;
   };
 
   /// This is a copy of module from parent node. All the reducer patterns will
   /// be applied to this instance.
-  ModuleOp module;
+  OwningOpRef<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;
 

mlir/lib/Reducer/ReductionNode.cpp

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())});
+  } else {
+    // Release the uninteresting module to save some memory.
+    module.release()->erase();
   }
 }
 
 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.
   ArrayRef<ReductionNode *> variantsFromParent =

mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp

 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
 // Parse and verify the input MLIR file.
-static LogicalResult loadModule(MLIRContext &context, OwningModuleRef &module,
+static LogicalResult loadModule(MLIRContext &context,
+                                OwningOpRef<ModuleOp> &module,
                                 StringRef inputFilename) {
   module = parseSourceFile(inputFilename, &context);
   if (!module)
     return failure();
 
   return success();
 }
 
LogicalResult mlir::mlirReduceMain(int argc, char **argv,
   }
 
   std::string errorMessage;
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output)
     return failure();
 
-  mlir::OwningModuleRef moduleRef;
+  OwningOpRef<ModuleOp> moduleRef;
   if (failed(loadModule(context, moduleRef, inputFilename)))
     return failure();
 
   auto errorHandler = [&](const Twine &msg) {
     return emitError(UnknownLoc::get(&context)) << msg;
   };
 
   // Reduction pass pipeline.
   PassManager pm(&context);
   if (failed(parser.addToPipeline(pm, errorHandler)))
     return failure();
 
-  ModuleOp m = moduleRef.get().clone();
+  OwningOpRef<ModuleOp> m = moduleRef.get().clone();
 
-  if (failed(pm.run(m)))
+  if (failed(pm.run(m.get())))
     return failure();
 
-  m.print(output->os());
+  m->print(output->os());
   output->keep();
 
   return success();
 }