[mlir][spirv] De-template deserialization

[antiagainst]

Summary

Previously for each op we generate a separate deserialization
method for it. Those deserialization methods duplicate the logic
of parsing operands/results/attributes and such.

This commit creates a generic method and let suitable op-specific
deserialization method to call into it.

wc -l SPIRVSerialization.inc: before 13290; after: 8304 (So -4986)

Depends On D93489

[antiagainst]

Fix comments

[hanchung]

mlir/lib/Target/SPIRV/Deserialization.cpp

2540–2542

drop trivial braces?

2558

[optional] I would move this to the beginning (ie size_t wordIndex = 0), and use wordIndex++ in the if (hasResult) section, so we can avoid magic numbers here.

mlir/lib/Target/SPIRV/Deserialization.cpp

 #include "llvm/ADT/bit.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "spirv-deserialization"
 
+//===----------------------------------------------------------------------===//
+// 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 StringRef decodeStringLiteral(ArrayRef<uint32_t> words,
                                             unsigned &wordIndex) {
   StringRef str(reinterpret_cast<const char *>(words.data() + wordIndex));
   wordIndex += str.size() / 4 + 1;
   return str;
 }
 
 /// Extracts the opcode from the given first word of a SPIR-V instruction.
 static inline spirv::Opcode extractOpcode(uint32_t word) {
   return static_cast<spirv::Opcode>(word & 0xffff);
 }
 
 /// Returns true if the given `block` is a function entry block.
 static inline bool isFnEntryBlock(Block *block) {
   return block->isEntryBlock() &&
          isa_and_nonnull<spirv::FuncOp>(block->getParentOp());
 }
 
 namespace {
+//===----------------------------------------------------------------------===//
+// Utility Definitions
+//===----------------------------------------------------------------------===//
+
 /// A struct for containing a header block's merge and continue targets.
 ///
 /// This struct is used to track original structured control flow info from
 /// SPIR-V blob. This info will be used to create spv.selection/spv.loop
 /// later.
 struct BlockMergeInfo {
   Block *mergeBlock;
   Block *continueBlock; // nullptr for spv.selection
struct DeferredStructTypeInfo {
 
   // The list of member types. For unresolved members, this list contains
   // place-holder empty types that will be updated later.
   SmallVector<Type, 4> memberTypes;
   SmallVector<spirv::StructType::OffsetInfo, 0> offsetInfo;
   SmallVector<spirv::StructType::MemberDecorationInfo, 0> memberDecorationsInfo;
 };
 
+//===----------------------------------------------------------------------===//
+// Deserializer Declaration
+//===----------------------------------------------------------------------===//
+
 /// A SPIR-V module serializer.
 ///
 /// A SPIR-V binary module is a single linear stream of instructions; each
 /// instruction is composed of 32-bit words. The first word of an instruction
 /// records the total number of words of that instruction using the 16
 /// higher-order bits. So this deserializer uses that to get instruction
 /// boundary and parse instructions and build a SPIR-V ModuleOp gradually.
 ///
private:
   /// Processing of Some instructions (like OpEntryPoint and OpExecutionMode)
   /// might need to be deferred, since they contain forward references to <id>s
   /// in the deserialized binary, but module in SPIR-V dialect expects these to
   /// be ssa-uses.
   LogicalResult processInstruction(spirv::Opcode opcode,
                                    ArrayRef<uint32_t> operands,
                                    bool deferInstructions = true);
 
+  /// Processes a SPIR-V instruction from the given `operands`. It should
+  /// deserialize into an op with the given `opName` and `numOperands`.
+  /// This method is a generic one for dispatching any SPIR-V ops without
+  /// variadic operands and attributes in TableGen definitions.
+  LogicalResult processOpWithoutGrammarAttr(ArrayRef<uint32_t> words,
+                                            StringRef opName, bool hasResult,
+                                            unsigned numOperands);
+
   /// Processes a OpUndef instruction. Adds a spv.Undef operation at the current
   /// insertion point.
   LogicalResult processUndef(ArrayRef<uint32_t> operands);
 
   LogicalResult processTypeForwardPointer(ArrayRef<uint32_t> operands);
 
   /// Method to dispatch to the specialized deserialization function for an
   /// operation in SPIR-V dialect that is a mirror of an instruction in the
private:
   /// instructions.
   llvm::SetVector<uint32_t> typeForwardPointerIDs;
 
   /// A list of all structs which have unresolved member types.
   SmallVector<DeferredStructTypeInfo, 0> deferredStructTypesInfos;
 };
 } // namespace
 
+//===----------------------------------------------------------------------===//
+// Deserializer Method Definitions
+//===----------------------------------------------------------------------===//
+
 Deserializer::Deserializer(ArrayRef<uint32_t> binary, MLIRContext *context)
     : binary(binary), context(context), unknownLoc(UnknownLoc::get(context)),
       module(createModuleOp()), opBuilder(module->body()) {}
 
 LogicalResult Deserializer::deserialize() {
   LLVM_DEBUG(llvm::dbgs() << "+++ starting deserialization +++\n");
 
   if (failed(processHeader()))
LogicalResult Deserializer::processInstruction(spirv::Opcode opcode,
   case spirv::Opcode::OpTypeForwardPointer:
     return processTypeForwardPointer(operands);
   default:
     break;
   }
   return dispatchToAutogenDeserialization(opcode, operands);
 }
 
+LogicalResult
+Deserializer::processOpWithoutGrammarAttr(ArrayRef<uint32_t> words,
+                                          StringRef opName, bool hasResult,
+                                          unsigned numOperands) {
+  SmallVector<Type, 1> resultTypes;
+  uint32_t valueID = 0;
+
+  size_t wordIndex= 0;
+  if (hasResult) {
+    if (wordIndex >= words.size())
+      return emitError(unknownLoc,
+                       "expected result type <id> while deserializing for ")
+             << opName;
+
+    // Decode the type <id>
+    auto type = getType(words[wordIndex]);
+    if (!type)
+      return emitError(unknownLoc, "unknown type result <id>: ")
+             << words[wordIndex];
+    resultTypes.push_back(type);
+    ++wordIndex;
+
+    // Decode the result <id>
+    if (wordIndex >= words.size())
+      return emitError(unknownLoc,
+                       "expected result <id> while deserializing for ")
+             << opName;
+    valueID = words[wordIndex];
+    ++wordIndex;
+  }
+
+  SmallVector<Value, 4> operands;
+  SmallVector<NamedAttribute, 4> attributes;
+
+  // Decode operands
+  size_t operandIndex = 0;
+  for (; operandIndex < numOperands && wordIndex < words.size();
+       ++operandIndex, ++wordIndex) {
+    auto arg = getValue(words[wordIndex]);
+    if (!arg)
+      return emitError(unknownLoc, "unknown result <id>: ") << words[wordIndex];
+    operands.push_back(arg);
+  }
+  if (operandIndex != numOperands) {
+    return emitError(
+               unknownLoc,
+               "found less operands than expected when deserializing for ")
+           << opName << "; only " << operandIndex << " of " << numOperands
+           << " processed";
+  }
+  if (wordIndex != words.size()) {
+    return emitError(
+               unknownLoc,
+               "found more operands than expected when deserializing for ")
+           << opName << "; only " << wordIndex << " of " << words.size()
+           << " processed";
+  }
+
+  // Attach attributes from decorations
+  if (decorations.count(valueID)) {
+    auto attrs = decorations[valueID].getAttrs();
+    attributes.append(attrs.begin(), attrs.end());
+  }
+
+  // Create the op and update bookkeeping maps
+  Location loc = createFileLineColLoc(opBuilder);
+  OperationState opState(loc, opName);
+  opState.addOperands(operands);
+  if (hasResult)
+    opState.addTypes(resultTypes);
+  opState.addAttributes(attributes);
+  Operation *op = opBuilder.createOperation(opState);
+  if (hasResult)
+    valueMap[valueID] = op->getResult(0);
+
+  if (op->hasTrait<OpTrait::IsTerminator>())
+    clearDebugLine();
+
+  return success();
+}
+
 LogicalResult Deserializer::processUndef(ArrayRef<uint32_t> operands) {
   if (operands.size() != 2) {
     return emitError(unknownLoc, "OpUndef instruction must have two operands");
   }
   auto type = getType(operands[0]);
   if (!type) {
     return emitError(unknownLoc, "unknown type <id> with OpUndef instruction");
   }
Deserializer::processOp<spirv::CopyMemoryOp>(ArrayRef<uint32_t> words) {
 
   return success();
 }
 
 // Pull in auto-generated Deserializer::dispatchToAutogenDeserialization() and
 // various Deserializer::processOp<...>() specializations.
 #define GET_DESERIALIZATION_FNS
 #include "mlir/Dialect/SPIRV/IR/SPIRVSerialization.inc"
+
 } // namespace
 
 namespace mlir {
 spirv::OwningSPIRVModuleRef spirv::deserialize(ArrayRef<uint32_t> binary,
                                                MLIRContext *context) {
   Deserializer deserializer(binary, context);
 
   if (failed(deserializer.deserialize()))
     return nullptr;
 
   return deserializer.collect();
 }
 } // namespace mlir

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

static void emitDecorationDeserialization(const Operator &op, StringRef tabs,
   }
 }
 
 /// Generates code to deserialize an SPV_Op `op` into `os`.
 static void emitDeserializationFunction(const Record *attrClass,
                                         const Record *record,
                                         const Operator &op, raw_ostream &os) {
   // If the record has 'autogenSerialization' set to 0, nothing to do
-  if (!record->getValueAsBit("autogenSerialization")) {
+  if (!record->getValueAsBit("autogenSerialization"))
     return;
-  }
+
   StringRef resultTypes("resultTypes"), valueID("valueID"), words("words"),
       wordIndex("wordIndex"), opVar("op"), operands("operands"),
       attributes("attributes");
+
+  // Method declaration
   os << formatv("template <> "
                 "LogicalResult\nDeserializer::processOp<{0}>(ArrayRef<"
                 "uint32_t> {1}) {{\n",
                 op.getQualCppClassName(), words);
+
+  // Special case for ops without attributes in TableGen definitions
+  if (op.getNumAttributes() == 0 && op.getNumVariableLengthOperands() == 0) {
+    os << formatv("  return processOpWithoutGrammarAttr("
+                  "{0}, \"{1}\", {2}, {3});\n}\n\n",
+                  words, op.getOperationName(),
+                  op.getNumResults() ? "true" : "false", op.getNumOperands());
+    return;
+  }
+
   os << formatv("  SmallVector<Type, 1> {0};\n", resultTypes);
   os << formatv("  size_t {0} = 0; (void){0};\n", wordIndex);
   os << formatv("  uint32_t {0} = 0; (void){0};\n", valueID);
 
   // Deserialize result information
   emitResultDeserialization(op, record->getLoc(), "  ", words, wordIndex,
                             resultTypes, valueID, os);
 
   os << formatv("  SmallVector<Value, 4> {0};\n", operands);
   os << formatv("  SmallVector<NamedAttribute, 4> {0};\n", attributes);
   // Operand deserialization
   emitOperandDeserialization(op, record->getLoc(), "  ", words, wordIndex,
                              operands, attributes, os);
 
+  // Decorations
+  emitDecorationDeserialization(op, "  ", valueID, attributes, os);
+
   os << formatv("  Location loc = createFileLineColLoc(opBuilder);\n");
   os << formatv("  auto {1} = opBuilder.create<{0}>(loc, {2}, {3}, {4}); "
                 "(void){1};\n",
                 op.getQualCppClassName(), opVar, resultTypes, operands,
                 attributes);
   if (op.getNumResults() == 1) {
     os << formatv("  valueMap[{0}] = {1}.getResult();\n\n", valueID, opVar);
   }
 
   // According to SPIR-V spec:
   // This location information applies to the instructions physically following
   // this instruction, up to the first occurrence of any of the following: the
   // next end of block.
   os << formatv("  if ({0}.hasTrait<OpTrait::IsTerminator>())\n", opVar);
   os << formatv("    clearDebugLine();\n");
-
-  // Decorations
-  emitDecorationDeserialization(op, "  ", valueID, attributes, os);
   os << "  return success();\n";
   os << "}\n\n";
 }
 
 /// Generates the prologue for the function that dispatches the deserialization
 /// based on the `opcode`.
 static void initDispatchDeserializationFn(StringRef opcode, StringRef words,
                                           raw_ostream &os) {