[mlir][spirv] NFC: Serialize all operands together if possible

[antiagainst]

Summary

SPIR-V ops can mix operands and attributes in the definition. These
operands and attributes are serialized in the exact order of the definition
to match SPIR-V binary format requirements. It can cause excessive
generated code bloat because we are emitting code to handle each
operand/attribute separately. We can probe first to check whether all
the operands are ahead of attributes. Then we can serialize all operands
together.

This removes ~1000 lines of code from the generated inc file.

Depends On D79444

[mravishankar]

Ah, thats nice. Simple change saves compilation time. Thanks!

[rriddle]

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

546

std::any_of?

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

 
 /// Generates code to serialize attributes of a SPV_Op `op` into `os`. The
 /// generates code extracts the attribute with name `attrName` from
 /// `operandList` of `op`.
 static void emitAttributeSerialization(const Attribute &attr,
                                        ArrayRef<SMLoc> loc, StringRef tabs,
                                        StringRef opVar, StringRef operandList,
                                        StringRef attrName, raw_ostream &os) {
-  os << tabs << formatv("auto attr = {0}.getAttr(\"{1}\");\n", opVar, attrName);
-  os << tabs << "if (attr) {\n";
+  os << tabs
+     << formatv("if (auto attr = {0}.getAttr(\"{1}\")) {{\n", opVar, attrName);
   if (attr.getAttrDefName() == "SPV_ScopeAttr" ||
       attr.getAttrDefName() == "SPV_MemorySemanticsAttr") {
     os << tabs
        << formatv("  {0}.push_back(prepareConstantInt({1}.getLoc(), "
                   "attr.cast<IntegerAttr>()));\n",
                   operandList, opVar);
   } else if (attr.getAttrDefName() == "I32ArrayAttr") {
     // Serialize all the elements of the array
static void emitAttributeSerialization(const Attribute &attr,
   }
   os << tabs << "}\n";
 }
 
 /// Generates code to serialize the operands of a SPV_Op `op` into `os`. The
 /// generated queries the SSA-ID if operand is a SSA-Value, or serializes the
 /// attributes. The `operands` vector is updated appropriately. `elidedAttrs`
 /// updated as well to include the serialized attributes.
-static void emitOperandSerialization(const Operator &op, ArrayRef<SMLoc> loc,
+static void emitArgumentSerialization(const Operator &op, ArrayRef<SMLoc> loc,
                                       StringRef tabs, StringRef opVar,
                                       StringRef operands, StringRef elidedAttrs,
                                       raw_ostream &os) {
+  using mlir::tblgen::Argument;
+
+  // SPIR-V ops can mix operands and attributes in the definition. These
+  // operands and attributes are serialized in the exact order of the definition
+  // to match SPIR-V binary format requirements. It can cause excessive
+  // generated code bloat because we are emitting code to handle each
+  // operand/attribute separately. So here we probe first to check whether all
+  // the operands are ahead of attributes. Then we can serialize all operands
+  // together.
+
+  // Whether all operands are ahead of all attributes in the op's spec.
+  bool areOperandsAheadOfAttrs = true;
+  // Find the first attribute.
+  const Argument *it = llvm::find_if(op.getArgs(), [](const Argument &arg) {
+    return arg.is<NamedAttribute *>();
+  });
+  // Check whether all following arguments are attributes.
+  for (const Argument *ie = op.arg_end(); it != ie; ++it) {
+    if (!it->is<NamedAttribute *>()) {
+      areOperandsAheadOfAttrs = false;
+      break;
+    }
+  }
+
+  // Serialize all operands together.
+  if (areOperandsAheadOfAttrs) {
+    if (op.getNumOperands() != 0) {
+      os << tabs
+         << formatv(
+                "for (Value operand : {0}.getOperation()->getOperands()) {{\n",
+                opVar);
+      os << tabs << "  auto id = getValueID(operand);\n";
+      os << tabs << "  assert(id && \"use before def!\");\n";
+      os << tabs << formatv("  {0}.push_back(id);\n", operands);
+      os << tabs << "}\n";
+    }
+    for (const NamedAttribute &attr : op.getAttributes()) {
+      emitAttributeSerialization(
+          (attr.attr.isOptional() ? attr.attr.getBaseAttr() : attr.attr), loc,
+          tabs, opVar, operands, attr.name, os);
+      os << tabs
+         << formatv("{0}.push_back(\"{1}\");\n", elidedAttrs, attr.name);
+    }
+    return;
+  }
+
+  // Serialize operands separately.
   auto operandNum = 0;
   for (unsigned i = 0, e = op.getNumArgs(); i < e; ++i) {
     auto argument = op.getArg(i);
     os << tabs << "{\n";
     if (argument.is<NamedTypeConstraint *>()) {
       os << tabs
          << formatv("  for (auto arg : {0}.getODSOperands({1})) {{\n", opVar,
                     operandNum);
       os << tabs << "    auto argID = getValueID(arg);\n";
       os << tabs << "    if (!argID) {\n";
       os << tabs
          << formatv("      return emitError({0}.getLoc(), "
                     "\"operand #{1} has a use before def\");\n",
                     opVar, operandNum);
       os << tabs << "    }\n";
       os << tabs << formatv("    {0}.push_back(argID);\n", operands);
       os << "    }\n";
       operandNum++;
     } else {
-      auto attr = argument.get<NamedAttribute *>();
+      NamedAttribute *attr = argument.get<NamedAttribute *>();
       auto newtabs = tabs.str() + "  ";
       emitAttributeSerialization(
           (attr->attr.isOptional() ? attr->attr.getBaseAttr() : attr->attr),
           loc, newtabs, opVar, operands, attr->name, os);
       os << newtabs
          << formatv("{0}.push_back(\"{1}\");\n", elidedAttrs, attr->name);
     }
     os << tabs << "}\n";
static void emitSerializationFunction(const Record *attrClass,
   // Serialize result information.
   if (op.getNumResults() == 1) {
     os << formatv("  uint32_t {0} = 0;\n", resultID);
     emitResultSerialization(op, record->getLoc(), "  ", opVar, operands,
                             resultID, os);
   }
 
   // Process arguments.
-  emitOperandSerialization(op, record->getLoc(), "  ", opVar, operands,
+  emitArgumentSerialization(op, record->getLoc(), "  ", opVar, operands,
                             elidedAttrs, os);
 
   if (record->isSubClassOf("SPV_ExtInstOp")) {
     os << formatv("  encodeExtensionInstruction({0}, \"{1}\", {2}, {3});\n",
                   opVar, record->getValueAsString("extendedInstSetName"),
                   record->getValueAsInt("extendedInstOpcode"), operands);
   } else {
     // Emit debug info.
     os << formatv("  emitDebugLine(functionBody, {0}.getLoc());\n", opVar);