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

3 files changed, 71 insertions, 38 deletions

diff --git a/mlir/lib/Support/CMakeLists.txt b/mlir/lib/Support/CMakeLists.txt
index 8f631626be3..7594a8c7bbe 100644
--- a/mlir/lib/Support/CMakeLists.txt
+++ b/mlir/lib/Support/CMakeLists.txt
@@ -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
diff --git a/mlir/lib/Support/MlirOptMain.cpp b/mlir/lib/Support/MlirOptMain.cpp
index 055692d64c4..9da2619c418 100644
--- a/mlir/lib/Support/MlirOptMain.cpp
+++ b/mlir/lib/Support/MlirOptMain.cpp
@@ -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);
diff --git a/mlir/lib/Support/ToolUtilities.cpp b/mlir/lib/Support/ToolUtilities.cpp
new file mode 100644
index 00000000000..60d0eee6b8a
--- /dev/null
+++ b/mlir/lib/Support/ToolUtilities.cpp
@@ -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);
+}