[Mlir-commits] [mlir] 9560f02 - [mlir] Add `enableSplitting` and `insertMarkerInOutput` options to `splitAndProcessBuffer`
River Riddle
llvmlistbot at llvm.org
Tue Jun 28 15:43:12 PDT 2022
Author: River Riddle
Date: 2022-06-28T15:42:35-07:00
New Revision: 9560f021410f6ecbf764129a59731b7f9777a882
URL: https://github.com/llvm/llvm-project/commit/9560f021410f6ecbf764129a59731b7f9777a882
DIFF: https://github.com/llvm/llvm-project/commit/9560f021410f6ecbf764129a59731b7f9777a882.diff
LOG: [mlir] Add `enableSplitting` and `insertMarkerInOutput` options to `splitAndProcessBuffer`
`enableSplitting` simply enables/disables whether we should split
or use the full buffer. `insertMarkerInOutput` toggles if split markers
should be inserted in between prcessed output chunks.
These options allow for merging the duplicate code paths we have
when splitting is optional.
Differential Revision: https://reviews.llvm.org/D128764
Added:
Modified:
mlir/include/mlir/Support/ToolUtilities.h
mlir/lib/Support/ToolUtilities.cpp
mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp
mlir/tools/mlir-pdll/mlir-pdll.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Support/ToolUtilities.h b/mlir/include/mlir/Support/ToolUtilities.h
index b2fcb012b3ee2..d2c89409c0653 100644
--- a/mlir/include/mlir/Support/ToolUtilities.h
+++ b/mlir/include/mlir/Support/ToolUtilities.h
@@ -32,10 +32,15 @@ using ChunkBufferHandler = function_ref<LogicalResult(
/// all results to `os`.
///
/// This is used to allow a large number of small independent tests to be put
-/// into a single file.
+/// into a single file. `enableSplitting` can be used to toggle if splitting
+/// should be enabled, e.g. to allow for merging split and non-split code paths.
+/// When `insertMarkerInOutput` is true, split markers (`//-----`) are placed
+/// between each of the processed output chunks.
LogicalResult
splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
- ChunkBufferHandler processChunkBuffer, raw_ostream &os);
+ ChunkBufferHandler processChunkBuffer, raw_ostream &os,
+ bool enableSplitting = true,
+ bool insertMarkerInOutput = false);
} // namespace mlir
#endif // MLIR_SUPPORT_TOOLUTILITIES_H
diff --git a/mlir/lib/Support/ToolUtilities.cpp b/mlir/lib/Support/ToolUtilities.cpp
index 7b52072a89767..0f9c4a933d008 100644
--- a/mlir/lib/Support/ToolUtilities.cpp
+++ b/mlir/lib/Support/ToolUtilities.cpp
@@ -21,7 +21,12 @@ using namespace mlir;
LogicalResult
mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
ChunkBufferHandler processChunkBuffer,
- raw_ostream &os) {
+ raw_ostream &os, bool enableSplitting,
+ bool insertMarkerInOutput) {
+ // If splitting is disabled, we process the full input buffer.
+ if (!enableSplitting)
+ return processChunkBuffer(std::move(originalBuffer), os);
+
const char splitMarkerConst[] = "// -----";
StringRef splitMarker(splitMarkerConst);
const int splitMarkerLen = splitMarker.size();
@@ -73,7 +78,7 @@ mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
// Process each chunk in turn.
bool hadFailure = false;
- for (auto &subBuffer : sourceBuffers) {
+ auto interleaveFn = [&](StringRef subBuffer) {
auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
@@ -82,7 +87,9 @@ mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
Twine(splitLine) + " offset ");
if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
hadFailure = true;
- }
+ };
+ llvm::interleave(sourceBuffers, os, interleaveFn,
+ insertMarkerInOutput ? "\n// -----\n" : "");
// If any fails, then return a failure of the tool.
return failure(hadFailure);
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index ccfc1a5ad0039..deffc892142f7 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -145,6 +145,7 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
// We use an explicit threadpool to avoid creating and joining/destroying
// threads for each of the split.
ThreadPool *threadPool = nullptr;
+
// Create a temporary context for the sake of checking if
// --mlir-disable-threading was passed on the command line.
// We use the thread-pool this context is creating, and avoid
@@ -153,23 +154,15 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
if (threadPoolCtx.isMultithreadingEnabled())
threadPool = &threadPoolCtx.getThreadPool();
- if (splitInputFile)
- return splitAndProcessBuffer(
- std::move(buffer),
- [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
- LogicalResult result = processBuffer(
- os, std::move(chunkBuffer), verifyDiagnostics, verifyPasses,
- allowUnregisteredDialects, preloadDialectsInContext,
- passManagerSetupFn, registry, threadPool);
- os << "// -----\n";
- return result;
- },
- outputStream);
-
- return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
- verifyPasses, allowUnregisteredDialects,
- preloadDialectsInContext, passManagerSetupFn, registry,
- threadPool);
+ auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer,
+ raw_ostream &os) {
+ return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
+ verifyPasses, allowUnregisteredDialects,
+ preloadDialectsInContext, passManagerSetupFn, registry,
+ threadPool);
+ };
+ return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
+ splitInputFile, /*insertMarkerInOutput=*/true);
}
LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
diff --git a/mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp b/mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp
index bce04e7959fe7..a8dc1df3e857f 100644
--- a/mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp
+++ b/mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp
@@ -98,13 +98,9 @@ LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
return sourceMgrHandler.verify();
};
- if (splitInputFile) {
- if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
- output->os())))
- return failure();
- } else if (failed(processBuffer(std::move(input), output->os()))) {
+ if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
+ output->os(), splitInputFile)))
return failure();
- }
output->keep();
return success();
diff --git a/mlir/tools/mlir-pdll/mlir-pdll.cpp b/mlir/tools/mlir-pdll/mlir-pdll.cpp
index f7f6f69105829..43e4aa5064240 100644
--- a/mlir/tools/mlir-pdll/mlir-pdll.cpp
+++ b/mlir/tools/mlir-pdll/mlir-pdll.cpp
@@ -186,13 +186,9 @@ int main(int argc, char **argv) {
return processBuffer(os, std::move(chunkBuffer), outputType, includeDirs,
dumpODS, includedFiles);
};
- if (splitInputFile) {
- if (failed(splitAndProcessBuffer(std::move(inputFile), processFn,
- outputStrOS)))
- return 1;
- } else if (failed(processFn(std::move(inputFile), outputStrOS))) {
+ if (failed(splitAndProcessBuffer(std::move(inputFile), processFn, outputStrOS,
+ splitInputFile)))
return 1;
- }
// Write the output.
bool shouldWriteOutput = true;
More information about the Mlir-commits
mailing list