[mlir] Execute all requested translations in MlirTranslateMain

[naibaf7]

Summary

Currently, MlirTranslateMain only executes one of the requested translations, and does not error if multiple are specified. This commit enables translations to be chained in the specified order.

This makes round-trip tests easier, since existing import/export passes can be reused and no combined round-trip passes have to be registered (example: mlir-translate -serialize-spirv -deserialize-spirv).

Additionally, by leveraging TranslateRegistration with file-to-file TranslateFunctions, generic pre- and post-processing can be added before/after conversion to/from MLIR.

[lattner]

This makes sense to me, thank you Fabian. That said, please also wait for another person to take a look at this. Some folks have opposed expanding the API of mlir-translate before and I'm not sure exactly what the boundaries we want are.

[naibaf7]

Add early-return on failure.

[naibaf7]

[mlir] Execute all requested translations in MlirTranslateMain

[Mogball]

LGTM. mlir-translate is intended first and foremost to be a testing tool. Making it easier to test with makes sense to me.

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

70

is this list guaranteed to be non-empty?

93

97

105

please spell out this auto

127

[naibaf7]

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

70

Yes, if it is empty (no option specified), then the ParseCommandLineOptions will error before this line is hit.

[naibaf7]

Address review.

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

 //===- MlirTranslateMain.cpp - MLIR Translation entry point ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Support/FileUtilities.h"
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/ToolUtilities.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
static llvm::cl::opt<bool> verifyDiagnostics(
       "verify-diagnostics",
       llvm::cl::desc("Check that emitted diagnostics match "
                      "expected-* lines on the corresponding line"),
       llvm::cl::init(false));
 
   llvm::InitLLVM y(argc, argv);
 
   // Add flags for all the registered translations.
-  llvm::cl::opt<const Translation *, false, TranslationParser>
-      translationRequested("", llvm::cl::desc("Translation to perform"),
+  llvm::cl::list<const Translation *, bool, TranslationParser>
+      translationsRequested("", llvm::cl::desc("Translations to perform"),
                             llvm::cl::Required);
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerTranslationCLOptions();
   llvm::cl::ParseCommandLineOptions(argc, argv, toolName);
 
   std::string errorMessage;
   std::unique_ptr<llvm::MemoryBuffer> input;
-  if (auto inputAlignment = translationRequested->getInputAlignment())
+  if (auto inputAlignment = translationsRequested[0]->getInputAlignment())
     input = openInputFile(inputFilename, *inputAlignment, &errorMessage);
   else
     input = openInputFile(inputFilename, &errorMessage);
   if (!input) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   // Processes the memory buffer with a new MLIRContext.
   auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
                            raw_ostream &os) {
+    // Temporary buffers for chained translation processing.
+    std::string dataIn;
+    std::string dataOut;
+    LogicalResult result = LogicalResult::success();
+
+    for (size_t i = 0, e = translationsRequested.size(); i < e; ++i) {
+      llvm::raw_ostream *stream;
+      llvm::raw_string_ostream dataStream(dataOut);
+
+      if (i == e - 1) {
+        // Output last translation to output.
+        stream = &os;
+      } else {
+        // Output translation to temporary data buffer.
+        stream = &dataStream;
+      }
+
+      const Translation *translationRequested = translationsRequested[i];
       MLIRContext context;
       context.allowUnregisteredDialects(allowUnregisteredDialects);
       context.printOpOnDiagnostic(!verifyDiagnostics);
       auto sourceMgr = std::make_shared<llvm::SourceMgr>();
       sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
-    if (!verifyDiagnostics) {
+      if (verifyDiagnostics) {
+        // 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);
+        result = sourceMgrHandler.verify();
+      } else {
         SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
-      return (*translationRequested)(sourceMgr, os, &context);
+        result = (*translationRequested)(sourceMgr, *stream, &context);
       }
+      if (failed(result))
+        return result;
 
-    // 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 (i < e - 1) {
+        // If there are further translations, create a new buffer with the
+        // output data.
+        dataIn = dataOut;
+        dataOut.clear();
+        ownedBuffer = llvm::MemoryBuffer::getMemBuffer(dataIn);
+      }
+    }
+    return result;
   };
 
   if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
                                    output->os(), splitInputFile)))
     return failure();
 
   output->keep();
   return success();
 }

mlir/test/Target/SPIRV/array-two-step-roundtrip.mlir

+// RUN: mlir-translate -no-implicit-module -split-input-file -serialize-spirv -deserialize-spirv %s | FileCheck %s
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+  spirv.func @array_stride(%arg0 : !spirv.ptr<!spirv.array<4x!spirv.array<4xf32, stride=4>, stride=128>, StorageBuffer>, %arg1 : i32, %arg2 : i32) "None" {
+    // CHECK: {{%.*}} = spirv.AccessChain {{%.*}}[{{%.*}}, {{%.*}}] : !spirv.ptr<!spirv.array<4 x !spirv.array<4 x f32, stride=4>, stride=128>, StorageBuffer>, i32, i32
+    %2 = spirv.AccessChain %arg0[%arg1, %arg2] : !spirv.ptr<!spirv.array<4x!spirv.array<4xf32, stride=4>, stride=128>, StorageBuffer>, i32, i32
+    spirv.Return
+  }
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+  // CHECK: spirv.GlobalVariable {{@.*}} : !spirv.ptr<!spirv.rtarray<f32, stride=4>, StorageBuffer>
+  spirv.GlobalVariable @var0 : !spirv.ptr<!spirv.rtarray<f32, stride=4>, StorageBuffer>
+  // CHECK: spirv.GlobalVariable {{@.*}} : !spirv.ptr<!spirv.rtarray<vector<4xf16>>, Input>
+  spirv.GlobalVariable @var1 : !spirv.ptr<!spirv.rtarray<vector<4xf16>>, Input>
+}