[MLIR] ODS TypeDefs: getChecked() and internal enhancements

[jdd]

Summary

Have the ODS TypeDef generator write the getChecked() definition.
Also, add to TypeParamCommaFormatter a JustParams format.

[rriddle]

mlir/tools/mlir-tblgen/TypeDefGen.cpp

88–90

Can you move the prependComma parameter to the end and have a default?

201

Please add comments for constant parameters, i.e. in this case /*prependComma=*/true

375

Same here and below.

498

Please wrap with braces here, the body is not trivial.

[jdd]

• Updates based on feedback fro River

mlir/test/lib/Dialect/Test/TestTypeDefs.td

def IntegerType : Test_Type<"TestInteger"> {
   let parser = [{
     if (parser.parseLess()) return Type();
     SignednessSemantics signedness;
     if (parseSignedness($_parser, signedness)) return mlir::Type();
     if ($_parser.parseComma()) return Type();
     int width;
     if ($_parser.parseInteger(width)) return Type();
     if ($_parser.parseGreater()) return Type();
-    return get(ctxt, signedness, width);
+    Location loc = $_parser.getEncodedSourceLoc($_parser.getNameLoc());
+    return getChecked(loc, signedness, width);
   }];
 
   // Any extra code one wants in the type's class declaration.
   let extraClassDeclaration = [{
     /// Signedness semantics.
     enum SignednessSemantics {
       Signless, /// No signedness semantics
       Signed,   /// Signed integer

mlir/test/mlir-tblgen/typedefs.td

let parameters = (
       "SomeCppStruct": $exampleCppType,
       ArrayRefParameter<"int", "Matrix dimensions">:$dims
   );
 
   let genVerifyInvariantsDecl = 1;
 
 // DECL-LABEL: class CompoundAType: public ::mlir::Type
 // DECL: static ::mlir::LogicalResult verifyConstructionInvariants(Location loc, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, SomeCppStruct exampleCppType, ::llvm::ArrayRef<int> dims);
-// DECL: static CompoundAType getChecked(Location loc, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, SomeCppStruct exampleCppType, ::llvm::ArrayRef<int> dims);
+// DECL: static ::mlir::Type getChecked(Location loc, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, SomeCppStruct exampleCppType, ::llvm::ArrayRef<int> dims);
 // DECL: static ::llvm::StringRef getMnemonic() { return "cmpnd_a"; }
 // DECL: static ::mlir::Type parse(::mlir::MLIRContext* ctxt, ::mlir::DialectAsmParser& parser);
 // DECL: void print(::mlir::DialectAsmPrinter& printer) const;
 // DECL: int getWidthOfSomething() const;
 // DECL: ::mlir::test::SimpleTypeA getExampleTdType() const;
 // DECL: SomeCppStruct getExampleCppType() const;
 }
 

mlir/tools/mlir-tblgen/TypeDefGen.cpp

 class TypeParamCommaFormatter : public llvm::detail::format_adapter {
 public:
   /// Choose the output format
   enum EmitFormat {
     /// Emit "parameter1Type parameter1Name, parameter2Type parameter2Name,
     /// [...]".
     TypeNamePairs,
 
-    /// Emit ", parameter1Type parameter1Name, parameter2Type parameter2Name,
-    /// [...]".
-    TypeNamePairsPrependComma,
-
     /// Emit "parameter1(parameter1), parameter2(parameter2), [...]".
-    TypeNameInitializer
+    TypeNameInitializer,
+
+    /// Emit "param1Name, param2Name, [...]".
+    JustParams,
   };
 
-  TypeParamCommaFormatter(EmitFormat emitFormat, ArrayRef<TypeParameter> params)
-      : emitFormat(emitFormat), params(params) {}
+  TypeParamCommaFormatter(EmitFormat emitFormat, ArrayRef<TypeParameter> params,
+                          bool prependComma = true)
+      : emitFormat(emitFormat), params(params), prependComma(prependComma) {}
 
   /// llvm::formatv will call this function when using an instance as a
   /// replacement value.
   void format(raw_ostream &os, StringRef options) override {
-    if (params.size() && emitFormat == EmitFormat::TypeNamePairsPrependComma)
+    if (params.size() && prependComma)
       os << ", ";
+
     switch (emitFormat) {
     case EmitFormat::TypeNamePairs:
-    case EmitFormat::TypeNamePairsPrependComma:
       interleaveComma(params, os,
                       [&](const TypeParameter &p) { emitTypeNamePair(p, os); });
       break;
     case EmitFormat::TypeNameInitializer:
       interleaveComma(params, os, [&](const TypeParameter &p) {
         emitTypeNameInitializer(p, os);
       });
       break;
+    case EmitFormat::JustParams:
+      interleaveComma(params, os,
+                      [&](const TypeParameter &p) { os << p.getName(); });
+      break;
     }
   }
 
 private:
   // Emit "paramType paramName".
   static void emitTypeNamePair(const TypeParameter &param, raw_ostream &os) {
     os << param.getCppType() << " " << param.getName();
   }
   // Emit "paramName(paramName)"
   void emitTypeNameInitializer(const TypeParameter &param, raw_ostream &os) {
     os << param.getName() << "(" << param.getName() << ")";
   }
 
   EmitFormat emitFormat;
   ArrayRef<TypeParameter> params;
+  bool prependComma;
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // GEN: TypeDef declarations
 //===----------------------------------------------------------------------===//
 

 static const char *const typeDefParsePrint = R"(
     static ::mlir::Type parse(::mlir::MLIRContext* ctxt, ::mlir::DialectAsmParser& parser);
     void print(::mlir::DialectAsmPrinter& printer) const;
 )";
 
 /// The code block for the verifyConstructionInvariants and getChecked.
 ///
 /// {0}: List of parameters, parameters style.
-/// {1}: C++ type class name.
 static const char *const typeDefDeclVerifyStr = R"(
     static ::mlir::LogicalResult verifyConstructionInvariants(Location loc{0});
-    static {1} getChecked(Location loc{0});
+    static ::mlir::Type getChecked(Location loc{0});
 )";
 
 /// Generate the declaration for the given typeDef class.
 static void emitTypeDefDecl(const TypeDef &typeDef, raw_ostream &os) {
   SmallVector<TypeParameter, 4> params;
   typeDef.getParameters(params);
 
   // Emit the beginning string template: either the singleton or parametric
   // template.
   if (typeDef.getNumParameters() == 0)
     os << formatv(typeDefDeclSingletonBeginStr, typeDef.getCppClassName(),
                   typeDef.getStorageNamespace(), typeDef.getStorageClassName());
   else
     os << formatv(typeDefDeclParametricBeginStr, typeDef.getCppClassName(),
                   typeDef.getStorageNamespace(), typeDef.getStorageClassName());
 
   // Emit the extra declarations first in case there's a type definition in
   // there.
   if (Optional<StringRef> extraDecl = typeDef.getExtraDecls())
     os << *extraDecl << "\n";
 
   TypeParamCommaFormatter emitTypeNamePairsAfterComma(
-      TypeParamCommaFormatter::EmitFormat::TypeNamePairsPrependComma, params);
+      TypeParamCommaFormatter::EmitFormat::TypeNamePairs, params);
   os << llvm::formatv("    static {0} get(::mlir::MLIRContext* ctxt{1});\n",
                       typeDef.getCppClassName(), emitTypeNamePairsAfterComma);
 
   // Emit the verify invariants declaration.
   if (typeDef.genVerifyInvariantsDecl())
-    os << llvm::formatv(typeDefDeclVerifyStr, emitTypeNamePairsAfterComma,
-                        typeDef.getCppClassName());
+    os << llvm::formatv(typeDefDeclVerifyStr, emitTypeNamePairsAfterComma);
 
   // Emit the mnenomic, if specified.
   if (auto mnenomic = typeDef.getMnemonic()) {
     os << "    static ::llvm::StringRef getMnemonic() { return \"" << mnenomic
        << "\"; }\n";
 
     // If mnemonic specified, emit print/parse declarations.
     os << typeDefParsePrint;

 /// {0}: storage class name.
 /// {1}: list of parameters.
 static const char *const typeDefStorageClassConstructorReturn = R"(
       return new (allocator.allocate<{0}>())
           {0}({1});
     }
 )";
 
+/// The code block for the getChecked definition.
+///
+/// {0}: List of parameters, parameters style.
+/// {1}: C++ type class name.
+/// {2}: Comma separated list of parameter names.
+static const char *const typeDefDefGetCheckeStr = R"(
+    ::mlir::Type {1}::getChecked(Location loc{0}) {{
+      return Base::getChecked(loc{2});
+    }
+)";
+
 /// Use tgfmt to emit custom allocation code for each parameter, if necessary.
 static void emitParameterAllocationCode(TypeDef &typeDef, raw_ostream &os) {
   SmallVector<TypeParameter, 4> parameters;
   typeDef.getParameters(parameters);
   auto fmtCtxt = FmtContext().addSubst("_allocator", "allocator");
   for (TypeParameter &parameter : parameters) {
     auto allocCode = parameter.getAllocator();
     if (allocCode) {
auto parameterNames = map_range(
       parameters, [](TypeParameter parameter) { return parameter.getName(); });
   auto parameterTypes = map_range(parameters, [](TypeParameter parameter) {
     return parameter.getCppType();
   });
   auto parameterList = join(parameterNames, ", ");
   auto parameterTypeList = join(parameterTypes, ", ");
 
   // 1) Emit most of the storage class up until the hashKey body.
-  os << formatv(
-      typeDefStorageClassBegin, typeDef.getStorageNamespace(),
+  os << formatv(typeDefStorageClassBegin, typeDef.getStorageNamespace(),
                 typeDef.getStorageClassName(),
                 TypeParamCommaFormatter(
-          TypeParamCommaFormatter::EmitFormat::TypeNamePairs, parameters),
+                    TypeParamCommaFormatter::EmitFormat::TypeNamePairs,
+                    parameters, /*prependComma=*/false),
                 TypeParamCommaFormatter(
-          TypeParamCommaFormatter::EmitFormat::TypeNameInitializer, parameters),
+                    TypeParamCommaFormatter::EmitFormat::TypeNameInitializer,
+                    parameters, /*prependComma=*/false),
                 parameterList, parameterTypeList);
 
   // 2) Emit the haskKey method.
   os << "  static ::llvm::hash_code hashKey(const KeyTy &key) {\n";
   // Extract each parameter from the key.
   for (size_t i = 0, e = parameters.size(); i < e; ++i)
-    os << formatv("      const auto &{0} = std::get<{1}>(key);\n",
+    os << llvm::formatv("      const auto &{0} = std::get<{1}>(key);\n",
                         parameters[i].getName(), i);
   // Then combine them all. This requires all the parameters types to have a
   // hash_value defined.
-  os << "      return ::llvm::hash_combine(";
-  interleaveComma(parameterNames, os);
-  os << ");\n";
-  os << "    }\n";
+  os << llvm::formatv(
+      "      return ::llvm::hash_combine({0});\n    }\n",
+      TypeParamCommaFormatter(TypeParamCommaFormatter::EmitFormat::JustParams,
+                              parameters, /* prependComma */ false));
 
   // 3) Emit the construct method.
   if (typeDef.hasStorageCustomConstructor())
     // If user wants to build the storage constructor themselves, declare it
     // here and then they can write the definition elsewhere.
     os << "    static " << typeDef.getStorageClassName()
        << " *construct(::mlir::TypeStorageAllocator &allocator, const KeyTy "
           "&key);\n";
static void emitTypeDefDef(TypeDef typeDef, raw_ostream &os) {
   typeDef.getParameters(parameters);
 
   // Emit the storage class, if requested and necessary.
   if (typeDef.genStorageClass() && typeDef.getNumParameters() > 0)
     emitStorageClass(typeDef, os);
 
   os << llvm::formatv(
       "{0} {0}::get(::mlir::MLIRContext* ctxt{1}) {{\n"
-      "  return Base::get(ctxt",
+      "  return Base::get(ctxt{2});\n}\n",
       typeDef.getCppClassName(),
       TypeParamCommaFormatter(
-          TypeParamCommaFormatter::EmitFormat::TypeNamePairsPrependComma,
+          TypeParamCommaFormatter::EmitFormat::TypeNamePairs, parameters),
+      TypeParamCommaFormatter(TypeParamCommaFormatter::EmitFormat::JustParams,
                               parameters));
-  for (TypeParameter &param : parameters)
-    os << ", " << param.getName();
-  os << ");\n}\n";
 
   // Emit the parameter accessors.
   if (typeDef.genAccessors())
     for (const TypeParameter &parameter : parameters) {
       SmallString<16> name = parameter.getName();
       name[0] = llvm::toUpper(name[0]);
       os << formatv("{0} {3}::get{1}() const { return getImpl()->{2}; }\n",
                     parameter.getCppType(), name, parameter.getName(),
                     typeDef.getCppClassName());
     }
 
+  // Generate getChecked() method.
+  if (typeDef.genVerifyInvariantsDecl()) {
+    os << llvm::formatv(
+        typeDefDefGetCheckeStr,
+        TypeParamCommaFormatter(
+            TypeParamCommaFormatter::EmitFormat::TypeNamePairs, parameters),
+        typeDef.getCppClassName(),
+        TypeParamCommaFormatter(TypeParamCommaFormatter::EmitFormat::JustParams,
+                                parameters));
+  }
+
   // If mnemonic is specified maybe print definitions for the parser and printer
   // code, if they're specified.
   if (typeDef.getMnemonic())
     emitParserPrinter(typeDef, os);
 }
 
 /// Emit the dialect printer/parser dispatcher. User's code should call these
 /// functions from their dialect's print/parse methods.