[mlir][OpDefGen] Add support for generating local functions for shared utilities

[rriddle]

Summary

This revision adds a new StaticVerifierFunctionEmitter class that emits local static functions in the .cpp file for shared operation verification. This class deduplicates shared operation verification code by emitting static functions alongside the op definitions. These methods are local to the definition file, and are invoked within the operation verify methods. The first bit of shared verification is for the type constraints used when verifying operands and results. An example is shown below:

static LogicalResult localVerify(...) {
  ...
}

LogicalResult OpA::verify(...) {
  if (failed(localVerify(...)))
    return failure();
  ...
}

LogicalResult OpB::verify(...) {
  if (failed(localVerify(...)))
    return failure();
  ...
}

This allowed for saving >400kb of code size from a downstream TensorFlow project (~15% of MLIR code size).

[jpienaar]

Nice

mlir/include/mlir/TableGen/Constraint.h

56

Is this only needed to be able to unique? Would exposing a compare function suffice to not need this? (To avoid exposing TableGen types from these types)

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

127

I don't follow the name. It doesn't really capture what this is doing - meaning, you have to read the comment and a bit of the code.

What about StaticVerifierFunctionEmitter ? local & ODS specific don't carry a much for me

216

Is .str() needed given the assignment to string?

[rriddle]

Resolve comments

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

127

I originally considered emitting other things aside from verification, but it's good to leave that to when we actually need it. Changed to your suggestion.

216

Yep.

[rriddle]

ping

[mehdi_amini]

I'd beef up the description to explain a bit more what this is doing.

[mehdi_amini]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

126

I think the description isn't super clear here either.

[jpienaar]

LG post Mehdi's suggestions

[rriddle]

Update

[rriddle]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

126

Added an example, which should make things a bit more clear.

mlir/include/mlir/TableGen/Constraint.h

public:
   // description is not provided, returns the TableGen def name.
   StringRef getDescription() const;
 
   // Constraint kind
   enum Kind { CK_Attr, CK_Region, CK_Successor, CK_Type, CK_Uncategorized };
 
   Kind getKind() const { return kind; }
 
+  /// Return the underlying tablegen definition of this constraint.
+  const llvm::Record *getDef() const { return def; }
+
 protected:
   Constraint(Kind kind, const llvm::Record *record);
 
   // The TableGen definition of this constraint.
   const llvm::Record *def;
 
 private:
   // What kind of constraint this is.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 static const char *const opCommentHeader = R"(
 //===----------------------------------------------------------------------===//
 // {0} {1}
 //===----------------------------------------------------------------------===//
 
 )";
 
 //===----------------------------------------------------------------------===//
+// LocalUtilityEmitter
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class handles emitting file local ODS specific utilities used within
+/// the definitions of generated operations.
+class LocalUtilityEmitter {
+public:
+  LocalUtilityEmitter(const llvm::RecordKeeper &records,
+                      ArrayRef<llvm::Record *> opDefs, raw_ostream &os,
+                      bool emitDecl);
+
+  /// Get the name of the local function used for the given type constraint.
+  /// These functions are used for operand and result constraints and have the
+  /// form:
+  ///   LogicalResult(Operation *op, Type type, StringRef valueKind,
+  ///                 unsigned valueGroupStartIndex);
+  StringRef getTypeConstraintFn(const llvm::Record *record) const {
+    auto it = localTypeConstraints.find(record);
+    assert(it != localTypeConstraints.end() && "expected valid constraint fn");
+    return it->second;
+  }
+
+private:
+  /// Returns a unique name to use when generating local methods.
+  static std::string getUniqueName(const llvm::RecordKeeper &records);
+
+  /// Emit local methods for the type constraints used within the provided op
+  /// definitions.
+  void emitTypeConstraintMethods(ArrayRef<llvm::Record *> opDefs,
+                                 raw_ostream &os, bool emitDecl);
+
+  /// A unique label for the file currently being generated. This is used to
+  /// ensure that the local functions have a unique name.
+  std::string uniqueOutputLabel;
+
+  /// A set of functions implementing type constraints, used for operand and
+  /// result verification.
+  llvm::DenseMap<const llvm::Record *, std::string> localTypeConstraints;
+};
+} // namespace
+
+LocalUtilityEmitter::LocalUtilityEmitter(const llvm::RecordKeeper &records,
+                                         ArrayRef<llvm::Record *> opDefs,
+                                         raw_ostream &os, bool emitDecl)
+    : uniqueOutputLabel(getUniqueName(records)) {
+  llvm::Optional<NamespaceEmitter> namespaceEmitter;
+  if (!emitDecl) {
+    os << formatv(opCommentHeader, "Local Utility Method", "Definitions");
+    namespaceEmitter.emplace(os, Operator(*opDefs[0]).getDialect());
+  }
+
+  emitTypeConstraintMethods(opDefs, os, emitDecl);
+}
+
+std::string
+LocalUtilityEmitter::getUniqueName(const llvm::RecordKeeper &records) {
+  // Use the input file name when generating a unique name.
+  std::string inputFilename = records.getInputFilename();
+
+  // Drop all but the base filename.
+  StringRef nameRef = llvm::sys::path::filename(inputFilename);
+  nameRef.consume_back(".td");
+
+  // Sanitize any invalid characters.
+  std::string uniqueName;
+  for (char c : nameRef) {
+    if (llvm::isAlnum(c) || c == '_')
+      uniqueName.push_back(c);
+    else
+      uniqueName.append(llvm::utohexstr((unsigned char)c));
+  }
+  return uniqueName;
+}
+
+void LocalUtilityEmitter::emitTypeConstraintMethods(
+    ArrayRef<llvm::Record *> opDefs, raw_ostream &os, bool emitDecl) {
+  // Collect a set of all of the used type constraints within the operation
+  // definitions.
+  llvm::SetVector<const llvm::Record *> typeConstraints;
+  for (Record *def : opDefs) {
+    Operator op(*def);
+    for (NamedTypeConstraint &operand : op.getOperands())
+      if (operand.hasPredicate())
+        typeConstraints.insert(operand.constraint.getDef());
+    for (NamedTypeConstraint &result : op.getResults())
+      if (result.hasPredicate())
+        typeConstraints.insert(result.constraint.getDef());
+  }
+
+  FmtContext fctx;
+  for (auto it : llvm::enumerate(typeConstraints)) {
+    // Generate an obscure and unique name for this type constraint.
+    std::string name = (Twine("__mlir_ods_local_type_constraint_") +
+                        uniqueOutputLabel + Twine(it.index()))
+                           .str();
+    localTypeConstraints.try_emplace(it.value(), name);
+
+    // Only generate the methods if we are generating definitions.
+    if (emitDecl)
+      continue;
+
+    TypeConstraint constraint(it.value());
+    os << "static ::mlir::LogicalResult " << name
+       << "(::mlir::Operation *op, ::mlir::Type type, ::llvm::StringRef "
+          "valueKind, unsigned valueGroupStartIndex) {\n";
+
+    os << "  if (!("
+       << tgfmt(constraint.getConditionTemplate(), &fctx.withSelf("type"))
+       << ")) {\n"
+       << formatv(
+              "    return op->emitOpError(valueKind) << \" #\" << "
+              "valueGroupStartIndex << \" must be {0}, but got \" << type;\n",
+              constraint.getDescription())
+       << "  }\n"
+       << "  return ::mlir::success();\n"
+       << "}\n\n";
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // Utility structs and functions
 //===----------------------------------------------------------------------===//
 
 // Replaces all occurrences of `match` in `str` with `substitute`.
 static std::string replaceAllSubstrs(std::string str, const std::string &match,
                                      const std::string &substitute) {
   std::string::size_type scanLoc = 0, matchLoc = std::string::npos;
   while ((matchLoc = str.find(match, scanLoc)) != std::string::npos) {

 //===----------------------------------------------------------------------===//
 // Op emitter
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Helper class to emit a record into the given output stream.
 class OpEmitter {
 public:
-  static void emitDecl(const Operator &op, raw_ostream &os);
-  static void emitDef(const Operator &op, raw_ostream &os);
+  static void emitDecl(const Operator &op, raw_ostream &os,
+                       const LocalUtilityEmitter &localUtilityEmitter);
+  static void emitDef(const Operator &op, raw_ostream &os,
+                      const LocalUtilityEmitter &localUtilityEmitter);
 
 private:
-  OpEmitter(const Operator &op);
+  OpEmitter(const Operator &op, const LocalUtilityEmitter &localUtilityEmitter);
 
   void emitDecl(raw_ostream &os);
   void emitDef(raw_ostream &os);
 
   // Generates the OpAsmOpInterface for this operation if possible.
   void genOpAsmInterface();
 
   // Generates the `getOperationName` method for this op.
private:
   // The wrapper operator class for querying information from this op.
   Operator op;
 
   // The C++ code builder for this op
   OpClass opClass;
 
   // The format context for verification code generation.
   FmtContext verifyCtx;
+
+  // The emitter containing all of the locally emitted utility functions.
+  const LocalUtilityEmitter &localUtilityEmitter;
 };
 } // end anonymous namespace
 
 // Populate the format context `ctx` with substitutions of attributes, operands
 // and results.
 // - attrGet corresponds to the name of the function to call to get value of
 //   attribute (the generated function call returns an Attribute);
 // - operandGet corresponds to the name of the function with which to retrieve
body << tgfmt("    if (!($0)) return $1\"attribute '$2' "
                   /*ctx=*/nullptr, tgfmt(condition, &ctx.withSelf(varName)),
                   emitErrorPrefix, attrName, attr.getDescription());
     if (allowMissingAttr)
       body << "  }\n";
     body << "  }\n";
   }
 }
 
-OpEmitter::OpEmitter(const Operator &op)
+OpEmitter::OpEmitter(const Operator &op,
+                     const LocalUtilityEmitter &localUtilityEmitter)
     : def(op.getDef()), op(op),
-      opClass(op.getCppClassName(), op.getExtraClassDeclaration()) {
+      opClass(op.getCppClassName(), op.getExtraClassDeclaration()),
+      localUtilityEmitter(localUtilityEmitter) {
   verifyCtx.withOp("(*this->getOperation())");
 
   genTraits();
 
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.
   genOpAsmInterface();
   genOpNameGetter();
OpEmitter::OpEmitter(const Operator &op,
   genCanonicalizerDecls();
   genFolderDecls();
   genOpInterfaceMethods();
   generateOpFormat(op, opClass);
   genSideEffectInterfaceMethods();
   genTypeInterfaceMethods();
 }
 
-void OpEmitter::emitDecl(const Operator &op, raw_ostream &os) {
-  OpEmitter(op).emitDecl(os);
+void OpEmitter::emitDecl(const Operator &op, raw_ostream &os,
+                         const LocalUtilityEmitter &localUtilityEmitter) {
+  OpEmitter(op, localUtilityEmitter).emitDecl(os);
 }
 
-void OpEmitter::emitDef(const Operator &op, raw_ostream &os) {
-  OpEmitter(op).emitDef(os);
+void OpEmitter::emitDef(const Operator &op, raw_ostream &os,
+                        const LocalUtilityEmitter &localUtilityEmitter) {
+  OpEmitter(op, localUtilityEmitter).emitDef(os);
 }
 
 void OpEmitter::emitDecl(raw_ostream &os) { opClass.writeDeclTo(os); }
 
 void OpEmitter::emitDef(raw_ostream &os) { opClass.writeDefTo(os); }
 
 void OpEmitter::genAttrGetters() {
   FmtContext fctx;
if (isOptional) {
                       "<< index << \" requires 0 or 1 element, but found \" << "
                       "valueGroup{0}.size();\n",
                       staticValue.index(), valueKind);
     }
 
     // Otherwise, if there is no predicate there is nothing left to do.
     if (!hasPredicate)
       continue;
-
     // Emit a loop to check all the dynamic values in the pack.
+    StringRef constraintFn = localUtilityEmitter.getTypeConstraintFn(
+        staticValue.value().constraint.getDef());
     body << "    for (::mlir::Value v : valueGroup" << staticValue.index()
-         << ") {\n";
-
-    auto constraint = staticValue.value().constraint;
-    body << "      (void)v;\n"
-         << "      if (!("
-         << tgfmt(constraint.getConditionTemplate(),
-                  &fctx.withSelf("v.getType()"))
-         << ")) {\n"
-         << formatv("        return emitOpError(\"{0} #\") << index "
-                    "<< \" must be {1}, but got \" << v.getType();\n",
-                    valueKind, constraint.getDescription())
-         << "      }\n" // if
+         << ") {\n"
+         << "      if (::mlir::failed(" << constraintFn
+         << "(getOperation(), v.getType(), \"" << valueKind << "\", index)))\n"
+         << "        return ::mlir::failure();\n"
          << "      ++index;\n"
-         << "    }\n"; // for
+         << "    }\n";
   }
 
   body << "  }\n";
 }
 
 void OpEmitter::genRegionVerifier(OpMethodBody &body) {
   // If we have no regions, there is nothing more to do.
   unsigned numRegions = op.getNumRegions();
void OpOperandAdaptorEmitter::emitDecl(const Operator &op, raw_ostream &os) {
   OpOperandAdaptorEmitter(op).adaptor.writeDeclTo(os);
 }
 
 void OpOperandAdaptorEmitter::emitDef(const Operator &op, raw_ostream &os) {
   OpOperandAdaptorEmitter(op).adaptor.writeDefTo(os);
 }
 
 // Emits the opcode enum and op classes.
-static void emitOpClasses(const std::vector<Record *> &defs, raw_ostream &os,
+static void emitOpClasses(const RecordKeeper &recordKeeper,
+                          const std::vector<Record *> &defs, raw_ostream &os,
                           bool emitDecl) {
   // First emit forward declaration for each class, this allows them to refer
   // to each others in traits for example.
   if (emitDecl) {
     os << "#if defined(GET_OP_CLASSES) || defined(GET_OP_FWD_DEFINES)\n";
     os << "#undef GET_OP_FWD_DEFINES\n";
     for (auto *def : defs) {
       Operator op(*def);
       NamespaceEmitter emitter(os, op.getDialect());
       os << "class " << op.getCppClassName() << ";\n";
     }
     os << "#endif\n\n";
   }
 
   IfDefScope scope("GET_OP_CLASSES", os);
+  if (defs.empty())
+    return;
+
+  // Generate all of the locally instantiated methods first.
+  LocalUtilityEmitter localUtilityEmitter(recordKeeper, defs, os, emitDecl);
   for (auto *def : defs) {
     Operator op(*def);
     NamespaceEmitter emitter(os, op.getDialect());
     if (emitDecl) {
       os << formatv(opCommentHeader, op.getQualCppClassName(), "declarations");
       OpOperandAdaptorEmitter::emitDecl(op, os);
-      OpEmitter::emitDecl(op, os);
+      OpEmitter::emitDecl(op, os, localUtilityEmitter);
     } else {
       os << formatv(opCommentHeader, op.getQualCppClassName(), "definitions");
       OpOperandAdaptorEmitter::emitDef(op, os);
-      OpEmitter::emitDef(op, os);
+      OpEmitter::emitDef(op, os, localUtilityEmitter);
     }
   }
 }
 
 // Emits a comma-separated list of the ops.
 static void emitOpList(const std::vector<Record *> &defs, raw_ostream &os) {
   IfDefScope scope("GET_OP_LIST", os);
 
getAllDerivedDefinitions(const RecordKeeper &recordKeeper,
 
   return defs;
 }
 
 static bool emitOpDecls(const RecordKeeper &recordKeeper, raw_ostream &os) {
   emitSourceFileHeader("Op Declarations", os);
 
   const auto &defs = getAllDerivedDefinitions(recordKeeper, "Op");
-  emitOpClasses(defs, os, /*emitDecl=*/true);
+  emitOpClasses(recordKeeper, defs, os, /*emitDecl=*/true);
 
   return false;
 }
 
 static bool emitOpDefs(const RecordKeeper &recordKeeper, raw_ostream &os) {
   emitSourceFileHeader("Op Definitions", os);
 
   const auto &defs = getAllDerivedDefinitions(recordKeeper, "Op");
   emitOpList(defs, os);
-  emitOpClasses(defs, os, /*emitDecl=*/false);
+  emitOpClasses(recordKeeper, defs, os, /*emitDecl=*/false);
 
   return false;
 }
 
 static mlir::GenRegistration
     genOpDecls("gen-op-decls", "Generate op declarations",
                [](const RecordKeeper &records, raw_ostream &os) {
                  return emitOpDecls(records, os);
                });
 
 static mlir::GenRegistration genOpDefs("gen-op-defs", "Generate op definitions",
                                        [](const RecordKeeper &records,
                                           raw_ostream &os) {
                                          return emitOpDefs(records, os);
                                        });