shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir][reducer] Add split-input-file to mlir-reduce (#184970)

Browse Source 
The tests for mlir-reduce are currently scattered. To centralize the
tests for mlir-reduce, I added the split-input-file feature to
mlir-reduce.It is part of
https://github.com/llvm/llvm-project/pull/184974.
This commit is contained in:
lonely eagle 2026-03-07 04:44:31 +08:00 committed by GitHub
parent ae4e71241f
commit a99d4a6311
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 43 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp
View File

@ -18,6 +18,7 @@
#include "mlir/Pass/PassManager.h"
#include "mlir/Reducer/Passes.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Support/ToolUtilities.h"
#include "mlir/Tools/ParseUtilities.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/SourceMgr.h"
@ -25,23 +26,6 @@
using namespace mlir;
// Parse and verify the input MLIR file. Returns null on error.
static OwningOpRef<Operation *> loadModule(MLIRContext &context,
StringRef inputFilename,
bool insertImplictModule) {
// Set up the input file.
std::string errorMessage;
auto file = openInputFile(inputFilename, &errorMessage);
if (!file) {
llvm::errs() << errorMessage << "\n";
return nullptr;
}
auto sourceMgr = std::make_shared<llvm::SourceMgr>();
sourceMgr->AddNewSourceBuffer(std::move(file), SMLoc());
return parseSourceFileForTool(sourceMgr, &context, insertImplictModule);
}
LogicalResult mlir::mlirReduceMain(int argc, char **argv,
MLIRContext &context) {
// Override the default '-h' and use the default PrintHelpMessage() which
@ -70,6 +54,18 @@ LogicalResult mlir::mlirReduceMain(int argc, char **argv,
llvm::cl::desc("Allow operation with no registered dialects"),
llvm::cl::init(false));
static llvm::cl::opt<std::string> splitInputFile(
"split-input-file", llvm::cl::ValueOptional,
llvm::cl::callback([&](const std::string &str) {
// Implicit value: use default marker if flag was used without
// value.
if (str.empty())
splitInputFile.setValue(kDefaultSplitMarker);
}),
llvm::cl::desc("Split the input file into chunks using the given or "
"default marker and process each chunk independently"),
llvm::cl::init(""));
llvm::cl::HideUnrelatedOptions(mlirReduceCategory);
llvm::InitLLVM y(argc, argv);
@ -93,27 +89,44 @@ LogicalResult mlir::mlirReduceMain(int argc, char **argv,
if (!output)
return failure();
OwningOpRef<Operation *> opRef =
loadModule(context, inputFilename, !noImplicitModule);
if (!opRef)
std::unique_ptr<llvm::MemoryBuffer> input =
openInputFile(inputFilename, &errorMessage);
if (!input) {
llvm::errs() << errorMessage << "\n";
return failure();
}
auto errorHandler = [&](const Twine &msg) {
return emitError(UnknownLoc::get(&context)) << msg;
};
// Reduction pass pipeline.
PassManager pm(&context, opRef.get()->getName().getStringRef());
if (failed(parser.addToPipeline(pm, errorHandler)))
return failure();
auto chunkFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
raw_ostream &os) {
auto sourceMgr = std::make_shared<llvm::SourceMgr>();
sourceMgr->AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());
OwningOpRef<Operation *> opRef =
parseSourceFileForTool(sourceMgr, &context, !noImplicitModule);
if (!opRef)
return failure();
// Reduction pass pipeline.
PassManager pm(&context, opRef.get()->getName().getStringRef());
if (failed(parser.addToPipeline(pm, errorHandler)))
return failure();
OwningOpRef<Operation *> op = opRef.get()->clone();
OwningOpRef<Operation *> op = opRef.get()->clone();
if (failed(pm.run(op.get())))
return failure();
if (failed(pm.run(op.get())))
return failure();
op.get()->print(output->os());
output->keep();
return success();
};
op.get()->print(output->os());
output->keep();
auto &splitInputFileDelimiter = splitInputFile.getValue();
if (!splitInputFileDelimiter.empty())
return splitAndProcessBuffer(std::move(input), chunkFn, output->os(),
splitInputFileDelimiter,
splitInputFileDelimiter);
return success();
return chunkFn(std::move(input), output->os());
}