[mlir][spirv] Add serialization control to emit symbol name

[antiagainst]

Summary

In SPIR-V, symbol names are encoded as OpName instructions.
They are not semantic impacting and can be omitted, which can
reduce the binary size.

mlir/include/mlir/Target/SPIRV/SPIRVBinaryUtils.h

 //===- SPIRVBinaryUtils.cpp - SPIR-V Binary Module Utils --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares common utilities for SPIR-V binary module.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_SPIRV_BINARY_UTILS_H_
 #define MLIR_TARGET_SPIRV_BINARY_UTILS_H_
 
-#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
-#include "mlir/Support/LogicalResult.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
+#include "mlir/Support/LLVM.h"
 
 #include <cstdint>
 
 namespace mlir {
 namespace spirv {
 
 /// SPIR-V binary header word count
 constexpr unsigned kHeaderWordCount = 5;
void appendModuleHeader(SmallVectorImpl<uint32_t> &header,
                         spirv::Version version, uint32_t idBound);
 
 /// Returns the word-count-prefixed opcode for an SPIR-V instruction.
 uint32_t getPrefixedOpcode(uint32_t wordCount, spirv::Opcode opcode);
 
 /// Encodes an SPIR-V `literal` string into the given `binary` vector.
 LogicalResult encodeStringLiteralInto(SmallVectorImpl<uint32_t> &binary,
                                       StringRef literal);
+
+/// Decodes a string literal in `words` starting at `wordIndex`. Update the
+/// latter to point to the position in words after the string literal.
+inline StringRef decodeStringLiteral(ArrayRef<uint32_t> words,
+                                     unsigned &wordIndex) {
+  StringRef str(reinterpret_cast<const char *>(words.data() + wordIndex));
+  wordIndex += str.size() / 4 + 1;
+  return str;
+}
+
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_TARGET_SPIRV_BINARY_UTILS_H_

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

 
 namespace mlir {
 struct LogicalResult;
 class MLIRContext;
 
 namespace spirv {
 class ModuleOp;
 
+struct SerializationOptions {
+  /// Whether to emit `OpName` instructions for SPIR-V symbol ops.
+  bool emitSymbolName = true;
+  /// Whether to emit `OpLine` location information for SPIR-V ops.
+  bool emitDebugInfo = false;
+};
+
 /// Serializes the given SPIR-V `module` 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,
-                        bool emitDebugInfo = false);
+                        const SerializationOptions &options = {});
 
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_TARGET_SPIRV_SERIALIZATION_H

mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "Deserializer.h"
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
+#include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "spirv-deserialization"
 

mlir/lib/Target/SPIRV/Deserialization/Deserializer.h

 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Builders.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include <cstdint>
 
-//===----------------------------------------------------------------------===//
-// Utility Functions
-//===----------------------------------------------------------------------===//
-
-/// Decodes a string literal in `words` starting at `wordIndex`. Update the
-/// latter to point to the position in words after the string literal.
-static inline llvm::StringRef
-decodeStringLiteral(llvm::ArrayRef<uint32_t> words, unsigned &wordIndex) {
-  llvm::StringRef str(reinterpret_cast<const char *>(words.data() + wordIndex));
-  wordIndex += str.size() / 4 + 1;
-  return str;
-}
-
 namespace mlir {
 namespace spirv {
 
 //===----------------------------------------------------------------------===//
 // Utility Definitions
 //===----------------------------------------------------------------------===//
 
 /// A struct for containing a header block's merge and continue targets.

mlir/lib/Target/SPIRV/Serialization/Serialization.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "spirv-serialization"
 
 namespace mlir {
 LogicalResult spirv::serialize(spirv::ModuleOp module,
                                SmallVectorImpl<uint32_t> &binary,
-                               bool emitDebugInfo) {
+                               const SerializationOptions &options) {
   if (!module.vce_triple().hasValue())
     return module.emitError(
         "module must have 'vce_triple' attribute to be serializeable");
 
-  Serializer serializer(module, emitDebugInfo);
+  Serializer serializer(module, options);
 
   if (failed(serializer.serialize()))
     return failure();
 
   LLVM_DEBUG(serializer.printValueIDMap(llvm::dbgs()));
 
   serializer.collect(binary);
   return success();
 }
 } // namespace mlir

mlir/lib/Target/SPIRV/Serialization/Serializer.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H
 #define MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/Target/SPIRV/Serialization.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace spirv {
 
 LogicalResult encodeInstructionInto(SmallVectorImpl<uint32_t> &binary,

 /// For the first word, the 16 high-order bits are the word count of the
 /// instruction, the 16 low-order bits are the opcode enumerant. The
 /// instructions then belong to different sections, which must be laid out in
 /// the particular order as specified in "2.4 Logical Layout of a Module" of
 /// the SPIR-V spec.
 class Serializer {
 public:
   /// Creates a serializer for the given SPIR-V `module`.
-  explicit Serializer(spirv::ModuleOp module, bool emitDebugInfo = false);
+  explicit Serializer(spirv::ModuleOp module,
+                      const SerializationOptions &options);
 
   /// Serializes the remembered SPIR-V module.
   LogicalResult serialize();
 
   /// Collects the final SPIR-V `binary`.
   void collect(SmallVectorImpl<uint32_t> &binary);
 
 #ifndef NDEBUG

 
 private:
   /// The SPIR-V module to be serialized.
   spirv::ModuleOp module;
 
   /// An MLIR builder for getting MLIR constructs.
   mlir::Builder mlirBuilder;
 
-  /// A flag which indicates if the debuginfo should be emitted.
-  bool emitDebugInfo = false;
+  /// Serialization options.
+  SerializationOptions options;
 
   /// A flag which indicates if the last processed instruction was a merge
   /// instruction.
   /// According to SPIR-V spec: "If a branch merge instruction is used, the last
   /// OpLine in the block must be before its merge instruction".
   bool lastProcessedWasMergeInst = false;
 
   /// The <id> of the OpString instruction, which specifies a file name, for

mlir/lib/Target/SPIRV/Serialization/Serializer.cpp

LogicalResult encodeInstructionInto(SmallVectorImpl<uint32_t> &binary,
                                     spirv::Opcode op,
                                     ArrayRef<uint32_t> operands) {
   uint32_t wordCount = 1 + operands.size();
   binary.push_back(spirv::getPrefixedOpcode(wordCount, op));
   binary.append(operands.begin(), operands.end());
   return success();
 }
 
-Serializer::Serializer(spirv::ModuleOp module, bool emitDebugInfo)
-    : module(module), mlirBuilder(module.getContext()),
-      emitDebugInfo(emitDebugInfo) {}
+Serializer::Serializer(spirv::ModuleOp module,
+                       const SerializationOptions &options)
+    : module(module), mlirBuilder(module.getContext()), options(options) {}
 
 LogicalResult Serializer::serialize() {
   LLVM_DEBUG(llvm::dbgs() << "+++ starting serialization +++\n");
 
   if (failed(module.verify()))
     return failure();
 
   // TODO: handle the other sections

 
 void Serializer::processCapability() {
   for (auto cap : module.vce_triple()->getCapabilities())
     (void)encodeInstructionInto(capabilities, spirv::Opcode::OpCapability,
                                 {static_cast<uint32_t>(cap)});
 }
 
 void Serializer::processDebugInfo() {
-  if (!emitDebugInfo)
+  if (!options.emitDebugInfo)
     return;
   auto fileLoc = module.getLoc().dyn_cast<FileLineColLoc>();
   auto fileName = fileLoc ? fileLoc.getFilename().strref() : "<unknown>";
   fileID = getNextID();
   SmallVector<uint32_t, 16> operands;
   operands.push_back(fileID);
   (void)spirv::encodeStringLiteralInto(operands, fileName);
   (void)encodeInstructionInto(debug, spirv::Opcode::OpString, operands);
LogicalResult Serializer::processDecoration(Location loc, uint32_t resultID,
   default:
     return emitError(loc, "unhandled decoration ") << decorationName;
   }
   return emitDecoration(resultID, decoration.getValue(), args);
 }
 
 LogicalResult Serializer::processName(uint32_t resultID, StringRef name) {
   assert(!name.empty() && "unexpected empty string for OpName");
+  if (!options.emitSymbolName)
+    return success();
 
   SmallVector<uint32_t, 4> nameOperands;
   nameOperands.push_back(resultID);
-  if (failed(spirv::encodeStringLiteralInto(nameOperands, name))) {
+  if (failed(spirv::encodeStringLiteralInto(nameOperands, name)))
     return failure();
-  }
   return encodeInstructionInto(names, spirv::Opcode::OpName, nameOperands);
 }
 
 template <>
 LogicalResult Serializer::processTypeDecoration<spirv::ArrayType>(
     Location loc, spirv::ArrayType type, uint32_t resultID) {
   if (unsigned stride = type.getArrayStride()) {
     // OpDecorate %arrayTypeSSA ArrayStride strideLiteral
LogicalResult Serializer::emitDecoration(uint32_t target,
   decorations.push_back(target);
   decorations.push_back(static_cast<uint32_t>(decoration));
   decorations.append(params.begin(), params.end());
   return success();
 }
 
 LogicalResult Serializer::emitDebugLine(SmallVectorImpl<uint32_t> &binary,
                                         Location loc) {
-  if (!emitDebugInfo)
+  if (!options.emitDebugInfo)
     return success();
 
   if (lastProcessedWasMergeInst) {
     lastProcessedWasMergeInst = false;
     return success();
   }
 
   auto fileLoc = loc.dyn_cast<FileLineColLoc>();
   if (fileLoc)
     (void)encodeInstructionInto(
         binary, spirv::Opcode::OpLine,
         {fileID, fileLoc.getLine(), fileLoc.getColumn()});
   return success();
 }
 } // namespace spirv
 } // namespace mlir

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

 
 // Deserializes the SPIR-V binary module stored in the file named as
 // `inputFilename` and returns a module containing the SPIR-V module.
 static OwningModuleRef deserializeModule(const llvm::MemoryBuffer *input,
                                          MLIRContext *context) {
   context->loadDialect<spirv::SPIRVDialect>();
 
   // Make sure the input stream can be treated as a stream of SPIR-V words
-  auto start = input->getBufferStart();
+  auto *start = input->getBufferStart();
   auto size = input->getBufferSize();
   if (size % sizeof(uint32_t) != 0) {
     emitError(UnknownLoc::get(context))
         << "SPIR-V binary module must contain integral number of 32-bit words";
     return {};
   }
 
   auto binary = llvm::makeArrayRef(reinterpret_cast<const uint32_t *>(start),
static LogicalResult serializeModule(ModuleOp module, raw_ostream &output) {
   module.walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
 
   if (spirvModules.empty())
     return module.emitError("found no 'spv.module' op");
 
   if (spirvModules.size() != 1)
     return module.emitError("found more than one 'spv.module' op");
 
-  if (failed(
-          spirv::serialize(spirvModules[0], binary, /*emitDebuginfo=*/false)))
+  if (failed(spirv::serialize(spirvModules[0], binary)))
     return failure();
 
   output.write(reinterpret_cast<char *>(binary.data()),
                binary.size() * sizeof(uint32_t));
 
   return mlir::success();
 }
 
static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
   auto spirvModules = srcModule.getOps<spirv::ModuleOp>();
 
   if (spirvModules.begin() == spirvModules.end())
     return srcModule.emitError("found no 'spv.module' op");
 
   if (std::next(spirvModules.begin()) != spirvModules.end())
     return srcModule.emitError("found more than one 'spv.module' op");
 
-  if (failed(spirv::serialize(*spirvModules.begin(), binary, emitDebugInfo)))
+  spirv::SerializationOptions options;
+  options.emitDebugInfo = emitDebugInfo;
+  if (failed(spirv::serialize(*spirvModules.begin(), binary, options)))
     return failure();
 
   MLIRContext deserializationContext(context->getDialectRegistry());
   // TODO: we should only load the required dialects instead of all dialects.
   deserializationContext.loadAllAvailableDialects();
   // Then deserialize to get back a SPIR-V module.
   OwningOpRef<spirv::ModuleOp> spirvModule =
       spirv::deserialize(binary, &deserializationContext);

mlir/unittests/Dialect/SPIRV/SerializationTest.cpp

 //===----------------------------------------------------------------------===//
 // Test Fixture
 //===----------------------------------------------------------------------===//
 
 class SerializationTest : public ::testing::Test {
 protected:
   SerializationTest() {
     context.getOrLoadDialect<mlir::spirv::SPIRVDialect>();
-    createModuleOp();
+    initModuleOp();
   }
 
-  void createModuleOp() {
+  /// Initializes an empty SPIR-V module op.
+  void initModuleOp() {
     OpBuilder builder(&context);
     OperationState state(UnknownLoc::get(&context),
                          spirv::ModuleOp::getOperationName());
     state.addAttribute("addressing_model",
                        builder.getI32IntegerAttr(static_cast<uint32_t>(
                            spirv::AddressingModel::Logical)));
     state.addAttribute("memory_model",
                        builder.getI32IntegerAttr(
                            static_cast<uint32_t>(spirv::MemoryModel::GLSL450)));
     state.addAttribute("vce_triple",
                        spirv::VerCapExtAttr::get(
                            spirv::Version::V_1_0, ArrayRef<spirv::Capability>(),
                            ArrayRef<spirv::Extension>(), &context));
     spirv::ModuleOp::build(builder, state);
     module = cast<spirv::ModuleOp>(Operation::create(state));
   }
 
-  Type getFloatStructType() {
-    OpBuilder opBuilder(module->getRegion());
-    llvm::SmallVector<Type, 1> elementTypes{opBuilder.getF32Type()};
+  /// Gets the `struct { float }` type.
+  spirv::StructType getFloatStructType() {
+    OpBuilder builder(module->getRegion());
+    llvm::SmallVector<Type, 1> elementTypes{builder.getF32Type()};
     llvm::SmallVector<spirv::StructType::OffsetInfo, 1> offsetInfo{0};
-    auto structType = spirv::StructType::get(elementTypes, offsetInfo);
-    return structType;
+    return spirv::StructType::get(elementTypes, offsetInfo);
   }
 
-  void addGlobalVar(Type type, llvm::StringRef name) {
-    OpBuilder opBuilder(module->getRegion());
+  /// Inserts a global variable of the given `type` and `name`.
+  spirv::GlobalVariableOp addGlobalVar(Type type, llvm::StringRef name) {
+    OpBuilder builder(module->getRegion());
     auto ptrType = spirv::PointerType::get(type, spirv::StorageClass::Uniform);
-    opBuilder.create<spirv::GlobalVariableOp>(
+    return builder.create<spirv::GlobalVariableOp>(
         UnknownLoc::get(&context), TypeAttr::get(ptrType),
-        opBuilder.getStringAttr(name), nullptr);
+        builder.getStringAttr(name), nullptr);
   }
 
+  /// Returns true if we can find a matching instruction in the SPIR-V blob.
   bool findInstruction(llvm::function_ref<bool(spirv::Opcode opcode,
                                                ArrayRef<uint32_t> operands)>
                            matchFn) {
     auto binarySize = binary.size();
-    auto begin = binary.begin();
+    auto *begin = binary.begin();
     auto currOffset = spirv::kHeaderWordCount;
 
     while (currOffset < binarySize) {
       auto wordCount = binary[currOffset] >> 16;
       if (!wordCount || (currOffset + wordCount > binarySize)) {
         return false;
       }
       spirv::Opcode opcode =
protected:
   OwningOpRef<spirv::ModuleOp> module;
   SmallVector<uint32_t, 0> binary;
 };
 
 //===----------------------------------------------------------------------===//
 // Block decoration
 //===----------------------------------------------------------------------===//
 
-TEST_F(SerializationTest, BlockDecorationTest) {
+TEST_F(SerializationTest, ContainsBlockDecoration) {
   auto structType = getFloatStructType();
   addGlobalVar(structType, "var0");
+
   ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary)));
+
   auto hasBlockDecoration = [](spirv::Opcode opcode,
                                ArrayRef<uint32_t> operands) -> bool {
     if (opcode != spirv::Opcode::OpDecorate || operands.size() != 2)
       return false;
     return operands[1] == static_cast<uint32_t>(spirv::Decoration::Block);
   };
   EXPECT_TRUE(findInstruction(hasBlockDecoration));
 }
+
+TEST_F(SerializationTest, ContainsSymbolName) {
+  auto structType = getFloatStructType();
+  addGlobalVar(structType, "var0");
+
+  spirv::SerializationOptions options;
+  options.emitSymbolName = true;
+  ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary, options)));
+
+  auto hasVarName = [](spirv::Opcode opcode, ArrayRef<uint32_t> operands) {
+    unsigned index = 1; // Skip the result <id>
+    return opcode == spirv::Opcode::OpName &&
+           spirv::decodeStringLiteral(operands, index) == "var0";
+  };
+  EXPECT_TRUE(findInstruction(hasVarName));
+}
+
+TEST_F(SerializationTest, DoesNotContainSymbolName) {
+  auto structType = getFloatStructType();
+  addGlobalVar(structType, "var0");
+
+  spirv::SerializationOptions options;
+  options.emitSymbolName = false;
+  ASSERT_TRUE(succeeded(spirv::serialize(module.get(), binary, options)));
+
+  auto hasVarName = [](spirv::Opcode opcode, ArrayRef<uint32_t> operands) {
+    unsigned index = 1; // Skip the result <id>
+    return opcode == spirv::Opcode::OpName &&
+           spirv::decodeStringLiteral(operands, index) == "var0";
+  };
+  EXPECT_FALSE(findInstruction(hasVarName));
+}