516ccce7fa
This allows to define custom splitters, which is interesting for non-MLIR inputs and outputs to `mlir-translate`. For example, one may use `; -----` as a splitter of `.ll` files. The splitters are now passed as arguments into `splitAndProcessBuffer`, the input splitter defaulting to the previous default (`// -----`) and the output splitter defaulting to the empty string, which also corresponds to the previous default. The behavior of the input split marker should not change at all; however, outputs now have one new line *more* than before if there is no splitter (old: `insertMarkerInOutput = false`, new: `outputSplitMarker = ""`) and one new line *less* if there is one. The value of the input splitter is exposed as a command line options of `mlir-translate` and other tools as an optional value to the previously existing flag `-split-input-file`, which defaults to the default splitter if not specified; the value of the output splitter is exposed with the new `-output-split-marker`, which default to the empty string in `mlir-translate` and the default splitter in the other tools. In short, the previous usage or omission of the flags should result in previous behavior (modulo the new lines mentioned before).
96 lines
3.8 KiB
C++
96 lines
3.8 KiB
C++
//===- ToolUtilities.cpp - MLIR Tool Utilities ----------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// 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"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
using namespace mlir;
|
|
|
|
LogicalResult
|
|
mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
|
|
ChunkBufferHandler processChunkBuffer,
|
|
raw_ostream &os, llvm::StringRef inputSplitMarker,
|
|
llvm::StringRef outputSplitMarker) {
|
|
// If splitting is disabled, we process the full input buffer.
|
|
if (inputSplitMarker.empty())
|
|
return processChunkBuffer(std::move(originalBuffer), os);
|
|
|
|
const int inputSplitMarkerLen = inputSplitMarker.size();
|
|
|
|
auto *origMemBuffer = originalBuffer.get();
|
|
SmallVector<StringRef, 8> rawSourceBuffers;
|
|
const int checkLen = 2;
|
|
// Split dropping the last checkLen chars to enable flagging near misses.
|
|
origMemBuffer->getBuffer().split(rawSourceBuffers,
|
|
inputSplitMarker.drop_back(checkLen));
|
|
if (rawSourceBuffers.empty())
|
|
return success();
|
|
|
|
// Add the original buffer to the source manager.
|
|
llvm::SourceMgr fileSourceMgr;
|
|
fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), SMLoc());
|
|
|
|
// Flag near misses by iterating over all the sub-buffers found when splitting
|
|
// with the prefix of the splitMarker. Use a sliding window where we only add
|
|
// a buffer as a sourceBuffer if terminated by a full match of the
|
|
// splitMarker, else flag a warning (if near miss) and extend the size of the
|
|
// buffer under consideration.
|
|
SmallVector<StringRef, 8> sourceBuffers;
|
|
StringRef prev;
|
|
for (auto buffer : rawSourceBuffers) {
|
|
if (prev.empty()) {
|
|
prev = buffer;
|
|
continue;
|
|
}
|
|
|
|
// Check that suffix is as expected and doesn't have any dash post.
|
|
bool expectedSuffix =
|
|
buffer.starts_with(inputSplitMarker.take_back(checkLen)) &&
|
|
buffer.size() > checkLen && buffer[checkLen] != '0';
|
|
if (expectedSuffix) {
|
|
sourceBuffers.push_back(prev);
|
|
prev = buffer.drop_front(checkLen);
|
|
} else {
|
|
// TODO: Consider making this a failure.
|
|
auto splitLoc = SMLoc::getFromPointer(buffer.data());
|
|
fileSourceMgr.PrintMessage(llvm::errs(), splitLoc,
|
|
llvm::SourceMgr::DK_Warning,
|
|
"near miss with file split marker");
|
|
prev = StringRef(prev.data(), prev.size() + inputSplitMarkerLen -
|
|
checkLen + buffer.size());
|
|
}
|
|
}
|
|
if (!prev.empty())
|
|
sourceBuffers.push_back(prev);
|
|
|
|
// Process each chunk in turn.
|
|
bool hadFailure = false;
|
|
auto interleaveFn = [&](StringRef subBuffer) {
|
|
auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
|
|
unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
|
|
auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
|
|
subBuffer, Twine("within split at ") +
|
|
origMemBuffer->getBufferIdentifier() + ":" +
|
|
Twine(splitLine) + " offset ");
|
|
if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
|
|
hadFailure = true;
|
|
};
|
|
llvm::interleave(sourceBuffers, os, interleaveFn,
|
|
(llvm::Twine(outputSplitMarker) + "\n").str());
|
|
|
|
// If any fails, then return a failure of the tool.
|
|
return failure(hadFailure);
|
|
}
|