Split out mlir-opt main into separate file.
authorJacques Pienaar <jpienaar@google.com>
Mon, 24 Jun 2019 15:41:52 +0000 (08:41 -0700)
committerjpienaar <jpienaar@google.com>
Mon, 24 Jun 2019 20:45:39 +0000 (13:45 -0700)
Enable reusing the real mlir-opt main from unit tests and in case where
additional initialization needs to happen before main is invoked (e.g., when
using different command line flag libraries).

PiperOrigin-RevId: 254764575

mlir/examples/Linalg/Linalg1/lib/CMakeLists.txt
mlir/include/mlir/Support/MlirOptMain.h [new file with mode: 0644]
mlir/lib/Support/CMakeLists.txt
mlir/lib/Support/MlirOptMain.cpp [new file with mode: 0644]
mlir/test/TestDialect/CMakeLists.txt
mlir/tools/mlir-opt/CMakeLists.txt
mlir/tools/mlir-opt/mlir-opt.cpp

index 67da27e..5cd3dd8 100644 (file)
@@ -69,5 +69,6 @@ target_link_libraries(linalg1-opt
   Linalg1DialectConstruction
   MLIRLLVMIR
   MLIRMlirOptLib
+  MLIROptMain
   ${LIBS}
   LLVMSupport)
diff --git a/mlir/include/mlir/Support/MlirOptMain.h b/mlir/include/mlir/Support/MlirOptMain.h
new file mode 100644 (file)
index 0000000..00a1e48
--- /dev/null
@@ -0,0 +1,38 @@
+//===- MlirOptMain.h - MLIR Optimizer Driver main ---------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// Main entry function for mlir-opt for when built as standalone binary.
+//
+//===----------------------------------------------------------------------===//
+
+#include <memory>
+#include <vector>
+
+namespace llvm {
+class raw_ostream;
+class MemoryBuffer;
+} // end namespace llvm
+namespace mlir {
+struct LogicalResult;
+class PassRegistryEntry;
+
+LogicalResult
+MlirOptMain(llvm::raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> buffer,
+            const std::vector<const PassRegistryEntry *> &passList,
+            bool splitInputFile, bool verifyDiagnostics, bool verifyPasses);
+
+} // end namespace mlir
index 84bbfeb..a82bd82 100644 (file)
@@ -1,5 +1,6 @@
 set(LLVM_OPTIONAL_SOURCES
   FileUtilities.cpp
+  MlirOptMain.cpp
   StorageUniquer.cpp
   TypeUtilities.cpp
 )
@@ -20,3 +21,11 @@ add_llvm_library(MLIRTypeUtilities
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
   )
 target_link_libraries(MLIRTypeUtilities MLIRIR)
+
+add_llvm_library(MLIROptMain
+  MlirOptMain.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
+  )
+target_link_libraries(MLIROptMain LLVMSupport)
diff --git a/mlir/lib/Support/MlirOptMain.cpp b/mlir/lib/Support/MlirOptMain.cpp
new file mode 100644 (file)
index 0000000..15b148a
--- /dev/null
@@ -0,0 +1,156 @@
+//===- MlirOptMain.cpp - MLIR Optimizer Driver ----------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This is a utility that runs an optimization pass and prints the result back
+// out. It is designed to support unit testing.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/MlirOptMain.h"
+#include "mlir/Analysis/Passes.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Module.h"
+#include "mlir/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/Regex.h"
+#include "llvm/Support/SourceMgr.h"
+
+using namespace mlir;
+using namespace llvm;
+using llvm::SMLoc;
+
+/// Perform the actions on the input file indicated by the command line flags
+/// within the specified context.
+///
+/// This typically parses the main source file, runs zero or more optimization
+/// passes, then prints the output.
+///
+static LogicalResult
+performActions(raw_ostream &os, bool verifyDiagnostics, bool verifyPasses,
+               SourceMgr &sourceMgr, MLIRContext *context,
+               const std::vector<const mlir::PassRegistryEntry *> &passList) {
+  std::unique_ptr<Module> module(parseSourceFile(sourceMgr, context));
+  if (!module)
+    return failure();
+
+  // Run each of the passes that were selected.
+  PassManager pm(verifyPasses);
+  for (const auto *passEntry : passList)
+    passEntry->addToPipeline(pm);
+
+  // Apply any pass manager command line options.
+  applyPassManagerCLOptions(pm);
+
+  // Run the pipeline.
+  if (failed(pm.run(module.get())))
+    return failure();
+
+  // Print the output.
+  module->print(os);
+  return success();
+}
+
+/// Parses the memory buffer.  If successfully, run a series of passes against
+/// it and print the result.
+static LogicalResult
+processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
+              bool verifyDiagnostics, bool verifyPasses,
+              const std::vector<const mlir::PassRegistryEntry *> &passList) {
+  // Tell sourceMgr about this buffer, which is what the parser will pick up.
+  SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
+
+  // Parse the input file.
+  MLIRContext context;
+
+  // If we are in verify diagnostics mode then we have a lot of work to do,
+  // otherwise just perform the actions without worrying about it.
+  if (!verifyDiagnostics) {
+    SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
+    return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
+                          &context, passList);
+  }
+
+  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
+
+  // Do any processing requested by command line flags.  We don't care whether
+  // these actions succeed or fail, we only care what diagnostics they produce
+  // and whether they match our expectations.
+  performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
+                 passList);
+
+  // Verify the diagnostic handler to make sure that each of the diagnostics
+  // matched.
+  return sourceMgrHandler.verify();
+}
+
+/// Split the specified file on a marker and process each chunk independently
+/// according to the normal processBuffer logic.  This is primarily used to
+/// allow a large number of small independent parser tests to be put into a
+/// single test, but could be used for other purposes as well.
+static LogicalResult splitAndProcessFile(
+    raw_ostream &os, std::unique_ptr<MemoryBuffer> originalBuffer,
+    bool verifyDiagnostics, bool verifyPasses,
+    const std::vector<const mlir::PassRegistryEntry *> &passList) {
+  const char marker[] = "// -----";
+  auto *origMemBuffer = originalBuffer.get();
+  SmallVector<StringRef, 8> sourceBuffers;
+  origMemBuffer->getBuffer().split(sourceBuffers, marker);
+
+  // Add the original buffer to the source manager.
+  SourceMgr fileSourceMgr;
+  fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), SMLoc());
+
+  bool hadUnexpectedResult = false;
+
+  // Process each chunk in turn.  If any fails, then return a failure of the
+  // tool.
+  for (auto &subBuffer : sourceBuffers) {
+    auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
+    unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
+    auto subMemBuffer = MemoryBuffer::getMemBufferCopy(
+        subBuffer, origMemBuffer->getBufferIdentifier() +
+                       Twine(" split at line #") + Twine(splitLine));
+    if (failed(processBuffer(os, std::move(subMemBuffer), verifyDiagnostics,
+                             verifyPasses, passList)))
+      hadUnexpectedResult = true;
+  }
+
+  return failure(hadUnexpectedResult);
+}
+
+LogicalResult
+mlir::MlirOptMain(raw_ostream &os, std::unique_ptr<MemoryBuffer> buffer,
+                  const std::vector<const mlir::PassRegistryEntry *> &passList,
+                  bool splitInputFile, bool verifyDiagnostics,
+                  bool verifyPasses) {
+  // The split-input-file mode is a very specific mode that slices the file
+  // up into small pieces and checks each independently.
+  if (splitInputFile)
+    return splitAndProcessFile(os, std::move(buffer), verifyDiagnostics,
+                               verifyPasses, passList);
+
+  return processBuffer(os, std::move(buffer), verifyDiagnostics, verifyPasses,
+                       passList);
+}
index 59fc920..dfd9686 100644 (file)
@@ -33,6 +33,7 @@ whole_archive_link(mlir-test-opt
 target_link_libraries(mlir-test-opt
   PRIVATE
   MLIRMlirOptLib
+  MLIROptMain
   MLIRTypeUtilities
   LLVMSupport
 )
index e9a85a7..bbca158 100644 (file)
@@ -11,6 +11,7 @@ set(LIB_LIBS
   MLIRSupport
 )
 add_llvm_library(MLIRMlirOptLib
+  main.cpp
   mlir-opt.cpp
 )
 target_link_libraries(MLIRMlirOptLib ${LIB_LIBS})
@@ -26,6 +27,7 @@ set(LIBS
   MLIRLinalg
   MLIRLLVMIR
   MLIRNVVMIR
+  MLIROptMain
   MLIRParser
   MLIRPass
   MLIRQuantizerTransforms
index bf300ec..35bba1f 100644 (file)
 // limitations under the License.
 // =============================================================================
 //
-// This is a command line utility that parses an MLIR file, runs an optimization
-// pass, then prints the result back out.  It is designed to support unit
-// testing.
+// Main entry function for mlir-opt for when built as standalone binary.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Passes.h"
-#include "mlir/IR/Attributes.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/Function.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/MLIRContext.h"
-#include "mlir/IR/Module.h"
-#include "mlir/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
-#include "mlir/Transforms/Passes.h"
+#include "mlir/Support/MlirOptMain.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/PrettyStackTrace.h"
-#include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
-using namespace mlir;
 using namespace llvm;
-using llvm::SMLoc;
+using namespace mlir;
 
 static cl::opt<std::string>
-inputFilename(cl::Positional, cl::desc("<input file>"), cl::init("-"));
+    inputFilename(cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
-static cl::opt<std::string>
-outputFilename("o", cl::desc("Output filename"), cl::value_desc("filename"),
-               cl::init("-"));
+static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
+                                           cl::value_desc("filename"),
+                                           cl::init("-"));
 
 static cl::opt<bool>
     splitInputFile("split-input-file",
@@ -69,106 +57,7 @@ static cl::opt<bool>
                  cl::desc("Run the verifier after each transformation pass"),
                  cl::init(true));
 
-static std::vector<const mlir::PassRegistryEntry *> *passList;
-
-enum OptResult { OptSuccess, OptFailure };
-
-/// Perform the actions on the input file indicated by the command line flags
-/// within the specified context.
-///
-/// This typically parses the main source file, runs zero or more optimization
-/// passes, then prints the output.
-///
-static OptResult performActions(SourceMgr &sourceMgr, MLIRContext *context) {
-  std::unique_ptr<Module> module(parseSourceFile(sourceMgr, context));
-  if (!module)
-    return OptFailure;
-
-  // Run each of the passes that were selected.
-  PassManager pm(verifyPasses);
-  for (const auto *passEntry : *passList)
-    passEntry->addToPipeline(pm);
-
-  // Apply any pass manager command line options.
-  applyPassManagerCLOptions(pm);
-
-  // Run the pipeline.
-  if (failed(pm.run(module.get())))
-    return OptFailure;
-
-  std::string errorMessage;
-  auto output = openOutputFile(outputFilename, &errorMessage);
-  if (!output) {
-    llvm::errs() << errorMessage << "\n";
-    exit(1);
-  }
-
-  // Print the output.
-  module->print(output->os());
-  output->keep();
-  return OptSuccess;
-}
-
-/// Parses the memory buffer.  If successfully, run a series of passes against
-/// it and print the result.
-static OptResult processFile(std::unique_ptr<MemoryBuffer> ownedBuffer) {
-  // Tell sourceMgr about this buffer, which is what the parser will pick up.
-  SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
-
-  // Parse the input file.
-  MLIRContext context;
-
-  // If we are in verify diagnostics mode then we have a lot of work to do,
-  // otherwise just perform the actions without worrying about it.
-  if (!verifyDiagnostics) {
-    SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
-    return performActions(sourceMgr, &context);
-  }
-
-  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
-
-  // Do any processing requested by command line flags.  We don't care whether
-  // these actions succeed or fail, we only care what diagnostics they produce
-  // and whether they match our expectations.
-  performActions(sourceMgr, &context);
-
-  // Verify the diagnostic handler to make sure that each of the diagnostics
-  // matched.
-  return failed(sourceMgrHandler.verify()) ? OptFailure : OptSuccess;
-}
-
-/// Split the specified file on a marker and process each chunk independently
-/// according to the normal processFile logic.  This is primarily used to
-/// allow a large number of small independent parser tests to be put into a
-/// single test, but could be used for other purposes as well.
-static OptResult
-splitAndProcessFile(std::unique_ptr<MemoryBuffer> originalBuffer) {
-  const char marker[] = "// -----";
-  auto *origMemBuffer = originalBuffer.get();
-  SmallVector<StringRef, 8> sourceBuffers;
-  origMemBuffer->getBuffer().split(sourceBuffers, marker);
-
-  // Add the original buffer to the source manager.
-  SourceMgr fileSourceMgr;
-  fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), SMLoc());
-
-  bool hadUnexpectedResult = false;
-
-  // Process each chunk in turn.  If any fails, then return a failure of the
-  // tool.
-  for (auto &subBuffer : sourceBuffers) {
-    auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
-    unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
-    auto subMemBuffer = MemoryBuffer::getMemBufferCopy(
-        subBuffer, origMemBuffer->getBufferIdentifier() +
-                       Twine(" split at line #") + Twine(splitLine));
-    if (processFile(std::move(subMemBuffer)))
-      hadUnexpectedResult = true;
-  }
-
-  return hadUnexpectedResult ? OptFailure : OptSuccess;
-}
+static std::vector<const PassRegistryEntry *> *passList;
 
 int main(int argc, char **argv) {
   llvm::PrettyStackTraceProgram x(argc, argv);
@@ -178,8 +67,8 @@ int main(int argc, char **argv) {
   registerPassManagerCLOptions();
 
   // Parse pass names in main to ensure static initialization completed.
-  llvm::cl::list<const mlir::PassRegistryEntry *, bool, PassNameParser>
-      passList("", llvm::cl::desc("Compiler passes to run"));
+  llvm::cl::list<const PassRegistryEntry *, bool, PassNameParser> passList(
+      "", llvm::cl::desc("Compiler passes to run"));
   ::passList = &passList;
   cl::ParseCommandLineOptions(argc, argv, "MLIR modular optimizer driver\n");
 
@@ -188,13 +77,15 @@ int main(int argc, char **argv) {
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
-    return OptFailure;
+    return 1;
   }
 
-  // The split-input-file mode is a very specific mode that slices the file
-  // up into small pieces and checks each independently.
-  if (splitInputFile)
-    return splitAndProcessFile(std::move(file));
+  auto output = openOutputFile(outputFilename, &errorMessage);
+  if (!output) {
+    llvm::errs() << errorMessage << "\n";
+    exit(1);
+  }
 
-  return processFile(std::move(file));
+  return failed(MlirOptMain(output->os(), std::move(file), passList,
+                            splitInputFile, verifyDiagnostics, verifyPasses));
 }