[mlir-reduce] Support parsing operations other than 'builtin.module' as top-level

[rkayaith]

Summary

This adds a --no-implicit-module option, which disables the insertion
of a top-level builtin.module during parsing. Although other ops can
now be parsed as top-level, the actual reduction passes are still
restricted to builtin.module as it didn't seem straightforward to
update them.

[rkayaith]

rebase

[rriddle]

mlir/lib/Reducer/ReductionTreePass.cpp

221–224

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

31–42

Can we change this function to just return an OwningOpRef<Operation *>? Using nullptr for failure?

[rkayaith]

address comments

[aartbik]

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

25

bazel file missing, again

[rkayaith]

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

25

is it expected that all MLIR changes should update the bazel files as well? i was under the impression that the bazel files were maintained separately by bazel users.

[aartbik]

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

25

No I don't think there is a general expectation for this. I think today I am the lucky bazel user that needs to keep this consistent ;-)

[rriddle]

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

25

There is no requirement to update bazel, it's a best effort done only by bazel users.

mlir/lib/Reducer/ReductionTreePass.cpp

LogicalResult ReductionTreePass::initialize(MLIRContext *context) {
   reducerPatterns = std::move(patterns);
   return success();
 }
 
 void ReductionTreePass::runOnOperation() {
   Operation *topOperation = getOperation();
   while (topOperation->getParentOp() != nullptr)
     topOperation = topOperation->getParentOp();
-  ModuleOp module = cast<ModuleOp>(topOperation);
+  ModuleOp module = dyn_cast<ModuleOp>(topOperation);
+  if (!module) {
+    emitError(getOperation()->getLoc())
+        << "top-level op must be 'builtin.module'";
+    return signalPassFailure();
+  }
 
   SmallVector<Operation *, 8> workList;
   workList.push_back(getOperation());
 
   do {
     Operation *op = workList.pop_back_val();
 
     for (Region &region : op->getRegions())

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

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Reducer/Passes.h"
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Tools/ParseUtilties.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
-// Parse and verify the input MLIR file.
-static LogicalResult loadModule(MLIRContext &context,
-                                OwningOpRef<ModuleOp> &module,
-                                StringRef inputFilename) {
-  module = parseSourceFile<ModuleOp>(inputFilename, &context);
-  if (!module)
-    return failure();
+// Parse and verify the input MLIR file. Returns null on error.
+OwningOpRef<Operation *> loadModule(MLIRContext &context,
+                                    StringRef inputFilename,
+                                    bool insertImplictModule) {
+  // Set up the input file.
+  std::string errorMessage;
+  auto file = openInputFile(inputFilename, &errorMessage);
+  if (!file) {
+    llvm::errs() << errorMessage << "\n";
+    return nullptr;
+  }
 
-  return success();
+  llvm::SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
+  return parseSourceFileForTool(sourceMgr, &context, insertImplictModule);
 }
 
 LogicalResult mlir::mlirReduceMain(int argc, char **argv,
                                    MLIRContext &context) {
   // Override the default '-h' and use the default PrintHelpMessage() which
   // won't print options in categories.
   static llvm::cl::opt<bool> help("h", llvm::cl::desc("Alias for -help"),
                                   llvm::cl::Hidden);
 
   static llvm::cl::OptionCategory mlirReduceCategory("mlir-reduce options");
 
   static llvm::cl::opt<std::string> inputFilename(
       llvm::cl::Positional, llvm::cl::desc("<input file>"),
       llvm::cl::cat(mlirReduceCategory));
 
   static llvm::cl::opt<std::string> outputFilename(
       "o", llvm::cl::desc("Output filename for the reduced test case"),
       llvm::cl::init("-"), llvm::cl::cat(mlirReduceCategory));
 
+  static llvm::cl::opt<bool> noImplicitModule{
+      "no-implicit-module",
+      llvm::cl::desc(
+          "Disable implicit addition of a top-level module op during parsing"),
+      llvm::cl::init(false)};
+
   llvm::cl::HideUnrelatedOptions(mlirReduceCategory);
 
   llvm::InitLLVM y(argc, argv);
 
   registerReducerPasses();
 
   PassPipelineCLParser parser("", "Reduction Passes to Run");
   llvm::cl::ParseCommandLineOptions(argc, argv,
                                     "MLIR test case reduction tool.\n");
 
   if (help) {
     llvm::cl::PrintHelpMessage();
     return success();
   }
 
   std::string errorMessage;
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output)
     return failure();
 
-  OwningOpRef<ModuleOp> moduleRef;
-  if (failed(loadModule(context, moduleRef, inputFilename)))
+  OwningOpRef<Operation *> opRef =
+      loadModule(context, inputFilename, !noImplicitModule);
+  if (!opRef)
     return failure();
 
   auto errorHandler = [&](const Twine &msg) {
     return emitError(UnknownLoc::get(&context)) << msg;
   };
 
   // Reduction pass pipeline.
-  PassManager pm(&context);
+  PassManager pm(&context, opRef.get()->getName().getStringRef());
   if (failed(parser.addToPipeline(pm, errorHandler)))
     return failure();
 
-  OwningOpRef<ModuleOp> m = moduleRef.get().clone();
+  OwningOpRef<Operation *> op = opRef.get()->clone();
 
-  if (failed(pm.run(m.get())))
+  if (failed(pm.run(op.get())))
     return failure();
 
-  m->print(output->os());
+  op.get()->print(output->os());
   output->keep();
 
   return success();
 }

mlir/test/mlir-reduce/invalid.mlir

+// UNSUPPORTED: system-windows
+// RUN: not mlir-reduce -opt-reduction-pass --no-implicit-module %s |& FileCheck %s --check-prefix=CHECK-PASS
+// RUN: not mlir-reduce -reduction-tree --no-implicit-module %s |& FileCheck %s --check-prefix=CHECK-TREE
+
+// The reduction passes are currently restricted to 'builtin.module'.
+// CHECK-PASS: error: Can't add pass '{{.+}}' restricted to 'builtin.module' on a PassManager intended to run on 'func.func'
+// CHECK-TREE: error: top-level op must be 'builtin.module'
+func.func private @foo()