Split out mlir-opt main into separate file.

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
This commit is contained in:
Jacques Pienaar 2019-06-24 08:41:52 -07:00 committed by jpienaar
parent 2ff1c01063
commit 257a654b72
7 changed files with 225 additions and 127 deletions

View file

@ -69,5 +69,6 @@ target_link_libraries(linalg1-opt
Linalg1DialectConstruction
MLIRLLVMIR
MLIRMlirOptLib
MLIROptMain
${LIBS}
LLVMSupport)

View file

@ -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

View 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)

View file

@ -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);
}

View file

@ -33,6 +33,7 @@ whole_archive_link(mlir-test-opt
target_link_libraries(mlir-test-opt
PRIVATE
MLIRMlirOptLib
MLIROptMain
MLIRTypeUtilities
LLVMSupport
)

View 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

View file

@ -15,42 +15,30 @@
// 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));
}