Let mlir-translate support -split-input-file

Similar to mlir-opt, having a -split-input-file mode is quite useful
in mlir-translate. It allows to put logically related tests in the
same test file for better organization.

PiperOrigin-RevId: 270805467
This commit is contained in:
Lei Zhang 2019-09-23 18:17:52 -07:00 committed by A. Unique TensorFlower
parent 75906bd565
commit 0e7edcfe7e
5 changed files with 145 additions and 41 deletions

View File

@ -0,0 +1,50 @@
//===- ToolUtilities.h - MLIR Tool Utilities --------------------*- 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.
// =============================================================================
//
// This file declares common utilities for implementing MLIR tools.
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_SUPPORT_TOOLUTILITIES_H
#define MLIR_SUPPORT_TOOLUTILITIES_H
#include "llvm/ADT/STLExtras.h"
#include <memory>
namespace llvm {
class MemoryBuffer;
}
namespace mlir {
class LogicalResult;
using ChunkBufferHandler = llvm::function_ref<LogicalResult(
std::unique_ptr<llvm::MemoryBuffer> chunkBuffer, llvm::raw_ostream &os)>;
/// Splits the specified buffer on a marker (`// -----`), processes each chunk
/// independently according to the normal `processChunkBuffer` logic, and writes
/// all results to `os`.
///
/// This is used to allow a large number of small independent tests to be put
/// into a single file.
LogicalResult
splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
ChunkBufferHandler processChunkBuffer,
llvm::raw_ostream &os);
} // namespace mlir
#endif // MLIR_SUPPORT_TOOLUTILITIES_H

View File

@ -3,12 +3,14 @@ set(LLVM_OPTIONAL_SOURCES
JitRunner.cpp
MlirOptMain.cpp
StorageUniquer.cpp
ToolUtilities.cpp
TranslateClParser.cpp
)
add_llvm_library(MLIRSupport
FileUtilities.cpp
StorageUniquer.cpp
ToolUtilities.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
@ -21,7 +23,11 @@ add_llvm_library(MLIROptMain
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
)
target_link_libraries(MLIROptMain LLVMSupport MLIRPass)
target_link_libraries(MLIROptMain
MLIRPass
LLVMSupport
MLIRSupport
)
add_llvm_library(MLIRTranslateClParser
TranslateClParser.cpp

View File

@ -30,6 +30,7 @@
#include "mlir/Parser.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/ToolUtilities.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/Regex.h"
@ -103,42 +104,6 @@ static LogicalResult processBuffer(raw_ostream &os,
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 PassPipelineCLParser &passPipeline) {
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, passPipeline)))
hadUnexpectedResult = true;
}
return failure(hadUnexpectedResult);
}
LogicalResult mlir::MlirOptMain(raw_ostream &os,
std::unique_ptr<MemoryBuffer> buffer,
const PassPipelineCLParser &passPipeline,
@ -147,8 +112,13 @@ LogicalResult mlir::MlirOptMain(raw_ostream &os,
// 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,
return splitAndProcessBuffer(
std::move(buffer),
[&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
verifyPasses, passPipeline);
},
os);
return processBuffer(os, std::move(buffer), verifyDiagnostics, verifyPasses,
passPipeline);

View File

@ -0,0 +1,57 @@
//===- ToolUtilities.cpp - MLIR Tool Utilities ----------------------------===//
//
// 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 file defines common utilities for implementing MLIR tools.
//
//===----------------------------------------------------------------------===//
#include "mlir/Support/ToolUtilities.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/Support/SourceMgr.h"
using namespace mlir;
LogicalResult
mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
ChunkBufferHandler processChunkBuffer,
raw_ostream &os) {
const char splitMarker[] = "// -----";
auto *origMemBuffer = originalBuffer.get();
SmallVector<StringRef, 8> sourceBuffers;
origMemBuffer->getBuffer().split(sourceBuffers, splitMarker);
// Add the original buffer to the source manager.
llvm::SourceMgr fileSourceMgr;
fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), llvm::SMLoc());
// Process each chunk in turn.
bool hadFailure = false;
for (auto &subBuffer : sourceBuffers) {
auto splitLoc = llvm::SMLoc::getFromPointer(subBuffer.data());
unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
subBuffer, origMemBuffer->getBufferIdentifier() +
Twine(" split at line #") + Twine(splitLine));
if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
hadFailure = true;
}
// If any fails, then return a failure of the tool.
return failure(hadFailure);
}

View File

@ -23,9 +23,11 @@
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Support/ToolUtilities.h"
#include "mlir/Support/TranslateClParser.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/ToolOutputFile.h"
using namespace mlir;
@ -38,6 +40,12 @@ static llvm::cl::opt<std::string>
outputFilename("o", llvm::cl::desc("Output filename"),
llvm::cl::value_desc("filename"), llvm::cl::init("-"));
static llvm::cl::opt<bool>
splitInputFile("split-input-file",
llvm::cl::desc("Split the input file into pieces and "
"process each chunk independently"),
llvm::cl::init(false));
int main(int argc, char **argv) {
llvm::InitLLVM y(argc, argv);
@ -45,6 +53,7 @@ int main(int argc, char **argv) {
llvm::cl::opt<const TranslateFunction *, false, TranslationParser>
translationRequested("", llvm::cl::desc("Translation to perform"),
llvm::cl::Required);
llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR translation driver\n");
std::string errorMessage;
@ -60,9 +69,21 @@ int main(int argc, char **argv) {
return 1;
}
/// Processes the memory buffer with a new MLIRContext.
auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
raw_ostream &os) {
MLIRContext context;
if (failed((*translationRequested)(std::move(input), output->os(), &context)))
return (*translationRequested)(std::move(ownedBuffer), os, &context);
};
if (splitInputFile) {
if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
output->os())))
return 1;
} else {
if (failed(processBuffer(std::move(input), output->os())))
return 1;
}
output->keep();
return 0;