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:
parent
2ff1c01063
commit
257a654b72
|
@ -69,5 +69,6 @@ target_link_libraries(linalg1-opt
|
|||
Linalg1DialectConstruction
|
||||
MLIRLLVMIR
|
||||
MLIRMlirOptLib
|
||||
MLIROptMain
|
||||
${LIBS}
|
||||
LLVMSupport)
|
||||
|
|
38
mlir/include/mlir/Support/MlirOptMain.h
Normal file
38
mlir/include/mlir/Support/MlirOptMain.h
Normal 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
|
|
@ -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)
|
||||
|
|
156
mlir/lib/Support/MlirOptMain.cpp
Normal file
156
mlir/lib/Support/MlirOptMain.cpp
Normal 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);
|
||||
}
|
|
@ -33,6 +33,7 @@ whole_archive_link(mlir-test-opt
|
|||
target_link_libraries(mlir-test-opt
|
||||
PRIVATE
|
||||
MLIRMlirOptLib
|
||||
MLIROptMain
|
||||
MLIRTypeUtilities
|
||||
LLVMSupport
|
||||
)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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));
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue