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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/184970.diff


1 Files Affected:

- (modified) mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp (+45-15) 


``````````diff
diff --git a/mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp b/mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp
index 34459b8564a25..61e968bc2063e 100644
--- a/mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp
+++ b/mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp
@@ -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"
@@ -70,6 +71,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 +106,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();
-
-  OwningOpRef<Operation *> op = opRef.get()->clone();
-
-  if (failed(pm.run(op.get())))
-    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();
+
+    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());
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/184970