[mlir] Add `enableSplitting` and `insertMarkerInOutput` options to `splitAndProcessBuffer`

[rriddle]

Summary

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.

[bzcheeseman]

Thanks River :)

mlir/lib/Support/ToolUtilities.cpp

91

would this be better done as an llvm::interleave call?

[rriddle]

mlir/lib/Support/ToolUtilities.cpp

91

Yeah, nice suggestion. Let me do that.

[lattner]

nice!!

mlir/include/mlir/Support/ToolUtilities.h

 using ChunkBufferHandler = function_ref<LogicalResult(
     std::unique_ptr<llvm::MemoryBuffer> chunkBuffer, 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.
+/// 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

mlir/lib/Support/ToolUtilities.cpp

 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 
 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();
 
   auto *origMemBuffer = originalBuffer.get();
   SmallVector<StringRef, 8> rawSourceBuffers;
   const int checkLen = 2;
   // Split dropping the last checkLen chars to enable flagging near misses.
if (expectedSuffix) {
                        prev.size() + splitMarkerLen - checkLen + buffer.size());
     }
   }
   if (!prev.empty())
     sourceBuffers.push_back(prev);
 
   // 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(
         subBuffer, Twine("within split at ") +
                        origMemBuffer->getBufferIdentifier() + ":" +
                        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);
 }

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
                                 bool preloadDialectsInContext) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // 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
   // creating any thread when disabled.
   MLIRContext threadPoolCtx;
   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,
+  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,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 const PassPipelineCLParser &passPipeline,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,

mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp

auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
     // In the diagnostic verification flow, we ignore whether the translation
     // failed (in most cases, it is expected to fail). Instead, we check if the
     // diagnostics were produced as expected.
     SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
     (void)(*translationRequested)(sourceMgr, os, &context);
     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()))) {
-    return failure();
-  }
+                                   output->os(), splitInputFile)))
 
   output->keep();
   return success();
 }

mlir/tools/mlir-pdll/mlir-pdll.cpp

int main(int argc, char **argv) {
   // up into small pieces and checks each independently.
   std::string outputStr;
   llvm::raw_string_ostream outputStrOS(outputStr);
   auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
                        raw_ostream &os) {
     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))) {
-    return 1;
-  }
+  if (failed(splitAndProcessBuffer(std::move(inputFile), processFn, outputStrOS,
+                                   splitInputFile)))
 
   // Write the output.
   bool shouldWriteOutput = true;
   if (writeIfChanged) {
     // Only update the real output file if there are any differences. This
     // prevents recompilation of all the files depending on it if there aren't
     // any.
     if (auto existingOrErr =