[mlir][spirv] Enable serializer to write SPIR-V modules into separate files (#152678)

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: output is written to a
single `os` that is passed around, and the 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

@@ -15,6 +15,7 @@
 
 #include "mlir/Support/LLVM.h"
 #include <cstdint>
+#include <string>
 
 namespace mlir {
 class MLIRContext;
@@ -27,12 +28,33 @@ 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`. This can either be a file prefix, or a relative or
+  /// or an absolute path followed by the prefix. For example:
+  ///
+  ///   * "foo" - Create files with a `foo` prefix in the current working
+  ///     directory. For example: `fooXYZ123`, `fooABC456` ... `fooXXXXXX`.
+  ///     The last 6 characters will be a unique combination as
+  ///     generated by `llvm::sys::fs::createUniqueFile`.
+  ///
+  ///   * "my/dir/foo" - Create files in `my/dir` with a `foo` prefix. The
+  ///     `my/dir` need to exists. For example: `fooXYZ123`, `fooABC456` ...
+  ///     `fooXXXXXX` will be created and stored in `/my/dir`. Filenames
+  ///     follow the same pattern as above.
+  ///
+  ///   * "/home/user/my/dir" - Same as above but using an absolute path.
+  std::string validationFilePrefix = "";
 };
 
-/// Serializes the given SPIR-V `module` and writes to `binary`. On failure,
+/// Serializes the given SPIR-V `moduleOp` and writes to `binary`. On failure,
 /// reports errors to the error handler registered with the MLIR context for
-/// `module`.
-LogicalResult serialize(ModuleOp module, SmallVectorImpl<uint32_t> &binary,
+/// `moduleOp`.
+LogicalResult serialize(ModuleOp moduleOp, SmallVectorImpl<uint32_t> &binary,
                         const SerializationOptions &options = {});
 
 } // namespace spirv

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

@@ -21,8 +21,11 @@
 #include "mlir/Target/SPIRV/Serialization.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
@@ -76,24 +79,66 @@ void registerFromSPIRVTranslation() {
 // Serialization registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult serializeModule(spirv::ModuleOp module,
-                                     raw_ostream &output) {
+static LogicalResult
+serializeModule(spirv::ModuleOp moduleOp, raw_ostream &output,
+                const spirv::SerializationOptions &options) {
   SmallVector<uint32_t, 0> binary;
-  if (failed(spirv::serialize(module, binary)))
+  if (failed(spirv::serialize(moduleOp, 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) {
+    size_t dirSeparator =
+        options.validationFilePrefix.find(llvm::sys::path::get_separator());
+    // If file prefix includes directory check if that directory exists.
+    if (dirSeparator != std::string::npos) {
+      llvm::StringRef parentDir =
+          llvm::sys::path::parent_path(options.validationFilePrefix);
+      if (!llvm::sys::fs::is_directory(parentDir))
+        return moduleOp.emitError(
+            "validation prefix directory does not exist\n");
+    }
+
+    SmallString<128> filename;
+    int fd = 0;
+
+    std::error_code errorCode = llvm::sys::fs::createUniqueFile(
+        options.validationFilePrefix + "%%%%%%", fd, filename);
+    if (errorCode)
+      return moduleOp.emitError("error creating validation output file: ")
+             << errorCode.message() << "\n";
+
+    llvm::raw_fd_ostream validationOutput(fd, /*shouldClose=*/true);
+    validationOutput.write(reinterpret_cast<char *>(binary.data()),
+                           sizeInBytes);
+    validationOutput.flush();
+  }
 
   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. Created files need to "
+          "be removed manually once processed."),
+      llvm::cl::init(""));
+
   TranslateFromMLIRRegistration toBinary(
       "serialize-spirv", "serialize SPIR-V dialect",
-      [](spirv::ModuleOp module, raw_ostream &output) {
-        return serializeModule(module, output);
+      [](spirv::ModuleOp moduleOp, raw_ostream &output) {
+        return serializeModule(moduleOp, output,
+                               {true, false, !validationFilesPrefix.empty(),
+                                validationFilesPrefix});
       },
       [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();

mlir/test/Target/SPIRV/mlir-translate.mlir

@@ -0,0 +1,29 @@
+// Check that `--spirv-save-validation-files-with-prefix` generates
+// a correct number of files.
+
+// REQUIRES: shell
+// RUN: rm -rf %t
+// RUN: mkdir %t && mlir-translate --serialize-spirv --no-implicit-module \
+// RUN: --split-input-file --spirv-save-validation-files-with-prefix=%t/foo %s \
+// RUN: && ls %t | wc -l | FileCheck %s
+// RUN: rm -rf %t
+
+// CHECK: 4
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+}