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

mlir/include/mlir/Support/ToolUtilities.h

@@ -0,0 +1,50 @@
+//===- ToolUtilities.h - MLIR Tool Utilities --------------------*- C++ -*-===//
+//
+// 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 declares common utilities for implementing MLIR tools.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_TOOLUTILITIES_H
+#define MLIR_SUPPORT_TOOLUTILITIES_H
+
+#include "llvm/ADT/STLExtras.h"
+#include <memory>
+
+namespace llvm {
+class MemoryBuffer;
+}
+
+namespace mlir {
+class LogicalResult;
+
+using ChunkBufferHandler = llvm::function_ref<LogicalResult(
+    std::unique_ptr<llvm::MemoryBuffer> chunkBuffer, llvm::raw_ostream &os)>;
+
+/// Splits the specified buffer on a marker (`// -----`), processes each chunk
+/// independently according to the normal `processChunkBuffer` logic, and writes
+/// all results to `os`.
+///
+/// This is used to allow a large number of small independent tests to be put
+/// into a single file.
+LogicalResult
+splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
+                      ChunkBufferHandler processChunkBuffer,
+                      llvm::raw_ostream &os);
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_TOOLUTILITIES_H

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

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);

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);
+}

mlir/tools/mlir-translate/mlir-translate.cpp

@@ -23,9 +23,11 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Support/ToolUtilities.h"
 #include "mlir/Support/TranslateClParser.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
@@ -38,6 +40,12 @@ static llvm::cl::opt<std::string>
     outputFilename("o", llvm::cl::desc("Output filename"),
                    llvm::cl::value_desc("filename"), llvm::cl::init("-"));
 
+static llvm::cl::opt<bool>
+    splitInputFile("split-input-file",
+                   llvm::cl::desc("Split the input file into pieces and "
+                                  "process each chunk independently"),
+                   llvm::cl::init(false));
+
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
 
@@ -45,6 +53,7 @@ int main(int argc, char **argv) {
   llvm::cl::opt<const TranslateFunction *, false, TranslationParser>
       translationRequested("", llvm::cl::desc("Translation to perform"),
                            llvm::cl::Required);
+
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR translation driver\n");
 
   std::string errorMessage;
@@ -60,9 +69,21 @@ int main(int argc, char **argv) {
     return 1;
   }
 
+  /// Processes the memory buffer with a new MLIRContext.
+  auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
+                           raw_ostream &os) {
     MLIRContext context;
-  if (failed((*translationRequested)(std::move(input), output->os(), &context)))
+    return (*translationRequested)(std::move(ownedBuffer), os, &context);
+  };
+
+  if (splitInputFile) {
+    if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
+                                     output->os())))
       return 1;
+  } else {
+    if (failed(processBuffer(std::move(input), output->os())))
+      return 1;
+  }
 
   output->keep();
   return 0;