[mlir][spirv] Enable serializer to write SPIR-V modules into separate files

By default, `mlir-translate` writes all output into a single file even
when `--split-input-file` is used. This is not an issue for text files
as they can be easily split with an output separator. However, this
causes issues with binary SPIR-V modules.

Firstly, a binary file with multiple modules is not a valid SPIR-V,
but will be created if multiple modules are specified in the same
file and separated by "// -----". This does not cause issues with
MLIR internal tools but does not work with SPIRV-Tools.

Secondly, splitting binary files after serialization is non-trivial,
when compared to text files, so using an external tool is not
desirable.

This patch adds a SPIR-V serialization option that write SPIR-V modules
to separate files in addition to writing them to the `mlir-translate`
output file. This is not the ideal solution and ideally `mlir-translate`
would allow generating multiple output files when `--split-input-file`
is used, however adding such functionality is again non-trival due to
how processing and splitting is done: function handles writing to
a single `os` that are passed around, and number of split buffers is not
known ahead of time. As such a I propose to have a SPIR-V internal option
that will dump modules to files in the form they can be processed by
`spirv-val`. The behaviour of the new added argument may be confusing,
but benefits from being internal to SPIR-V target.

Alternatively, we could expose the spirv option in `mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp`,
and slice the output file on the SPIR-V magic number, and not keep
the file generated by default by `mlir-translate`. This would be a bit
cleaner in API sense, as it would not generate the additional file
containing all modules together. However, it pushes SPIR-V specific
code into the generic part of the `mlir-translate` and slicing is
potentially more error prone that just writing a single module
after it was serialized.

Author: IgWod-IMG

mlir/include/mlir/Target/SPIRV/Serialization.h

@@ -27,6 +27,14 @@ struct SerializationOptions {
   bool emitSymbolName = true;
   /// Whether to emit `OpLine` location information for SPIR-V ops.
   bool emitDebugInfo = false;
+  /// Whether to store a module to an additional file during
+  /// serialization. This is used to store the SPIR-V module to the
+  /// file in addition to writing it to `os` passed from the calling
+  /// tool. This saved file is later used for validation.
+  bool saveModuleForValidation = false;
+  /// A prefix prepended to the file used when `saveModuleForValidation`
+  /// is set to `true`.
+  std::string validationFilePrefix = "";
 };
 
 /// Serializes the given SPIR-V `module` and writes to `binary`. On failure,

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

@@ -23,6 +23,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
@@ -76,24 +77,57 @@ void registerFromSPIRVTranslation() {
 // Serialization registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult serializeModule(spirv::ModuleOp module,
-                                     raw_ostream &output) {
+// Static variable is probably not ideal, but it lets us have unique files names
+// without taking additional parameters from `mlir-translate`.
+static size_t validationFileCounter = 0;
+
+static LogicalResult
+serializeModule(spirv::ModuleOp module, raw_ostream &output,
+                const spirv::SerializationOptions &options) {
+
   SmallVector<uint32_t, 0> binary;
   if (failed(spirv::serialize(module, binary)))
     return failure();
 
-  output.write(reinterpret_cast<char *>(binary.data()),
-               binary.size() * sizeof(uint32_t));
+  size_t sizeInBytes = binary.size() * sizeof(uint32_t);
+
+  output.write(reinterpret_cast<char *>(binary.data()), sizeInBytes);
+
+  if (options.saveModuleForValidation) {
+    std::string errorMessage;
+    std::string filename =
+        options.validationFilePrefix + std::to_string(validationFileCounter++);
+    auto validationOutput = openOutputFile(filename, &errorMessage);
+    if (!validationOutput) {
+      llvm::errs() << errorMessage << "\n";
+      return failure();
+    }
+    validationOutput->os().write(reinterpret_cast<char *>(binary.data()),
+                                 sizeInBytes);
+    validationOutput->keep();
+  }
 
   return mlir::success();
 }
 
 namespace mlir {
 void registerToSPIRVTranslation() {
+  static llvm::cl::opt<std::string> validationFilesPrefix(
+      "spirv-save-validation-files-with-prefix",
+      llvm::cl::desc(
+          "When non-empty string is passed each serialized SPIR-V module is "
+          "saved to an additional file that starts with the given prefix. This "
+          "is used to generate separate binaries for validation, where "
+          "`--split-input-file` normally combines all outputs into one. The "
+          "one combined output (`-o`) is still written."),
+      llvm::cl::init(""));
+
   TranslateFromMLIRRegistration toBinary(
       "serialize-spirv", "serialize SPIR-V dialect",
       [](spirv::ModuleOp module, raw_ostream &output) {
-        return serializeModule(module, output);
+        return serializeModule(module, output,
+                               {true, false, (validationFilesPrefix != ""),
+                                validationFilesPrefix});
       },
       [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();