[MLIR][TableGen] Add genMnemonicAlias field for OpAsm{Type,Attr}Interface

Author: ZenithalHourlyRate

mlir/docs/DefiningDialects/AttributesAndTypes.md

@@ -105,6 +105,9 @@ def My_IntegerType : MyDialect_Type<"Integer", "int"> {
 
   /// Indicate that our type will add additional verification to the parameters.
   let genVerifyDecl = 1;
+
+  /// Indicate that our type will use the mnemonic as alias in assembly.
+  let genMnemonicAlias = 1;
 }
 ```
 
@@ -160,6 +163,9 @@ def My_IntegerAttr : MyDialect_Attr<"Integer", "int"> {
   /// Indicate to the ODS generator that we do not want the default builders,
   /// as we have defined our own simpler ones.
   let skipDefaultBuilders = 1;
+
+  /// Indicate that our attribute will use the mnemonic as alias in assembly.
+  let genMnemonicAlias = 1;
 }
 ```
 
@@ -1189,6 +1195,13 @@ Note that these are mechanisms intended for long-tail cases by power users; for
 not-yet-implemented widely-applicable cases, improving the infrastructure is
 preferable.
 
+### Mnemonic Alias in Assembly
+
+Attribute and Type can use aliases in the assembly to reduce verbosity.
+In such cases, `OpAsmAttrInterface` and `OpAsmTypeInterface` can be used to generate aliases.
+Often, a simple mnemonic alias is enough; then enabling `genMnemonicAlias` automatically
+generates an `getAlias` implementation using the Attribute or Type's mnemonic.
+
 ### Registering with the Dialect
 
 Once the attributes and types have been defined, they must then be registered

mlir/include/mlir/IR/AttrTypeBase.td

@@ -249,6 +249,9 @@ class AttrOrTypeDef<string valueType, string name, list<Trait> defTraits,
   // generated code is placed inside the class's C++ namespace. `$cppClass` is
   // replaced by the class name.
   code extraClassDefinition = [{}];
+
+  // Generate a default 'getAlias' method for OpAsm{Type,Attr}Interface.
+  bit genMnemonicAlias = 0;
 }
 
 // Define a new attribute, named `name`, belonging to `dialect` that inherits

mlir/include/mlir/TableGen/AttrOrTypeDef.h

@@ -216,6 +216,10 @@ class AttrOrTypeDef {
   /// Returns the def's extra class definition code.
   std::optional<StringRef> getExtraDefs() const;
 
+  /// Returns true if we need to generate a default 'getAlias' implementation
+  /// using the mnemonic.
+  bool genMnemonicAlias() const;
+
   /// Get the code location (for error printing).
   ArrayRef<SMLoc> getLoc() const;
 

mlir/lib/TableGen/AttrOrTypeDef.cpp

@@ -205,6 +205,10 @@ std::optional<StringRef> AttrOrTypeDef::getExtraDefs() const {
   return value.empty() ? std::optional<StringRef>() : value;
 }
 
+bool AttrOrTypeDef::genMnemonicAlias() const {
+  return def->getValueAsBit("genMnemonicAlias");
+}
+
 ArrayRef<SMLoc> AttrOrTypeDef::getLoc() const { return def->getLoc(); }
 
 bool AttrOrTypeDef::skipDefaultBuilders() const {

mlir/test/IR/op-asm-interface.mlir

@@ -71,6 +71,16 @@ func.func @alias_from_op_asm_type_interface() {
 
 // -----
 
+// CHECK: !op_asm_type_interface_tablegen_default =
+!type = !test.op_asm_type_interface_tablegen_default
+
+func.func @alias_from_op_asm_type_interface_tablegen_default() {
+  %0 = "test.result_name_from_type"() : () -> !type
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // Test OpAsmAttrInterface
 //===----------------------------------------------------------------------===//
@@ -82,3 +92,13 @@ func.func @test_op_asm_attr_interface() {
   %1 = "test.result_name_from_type"() {attr = #attr} : () -> !test.op_asm_type_interface
   return
 }
+
+// -----
+
+// CHECK: #op_asm_attr_interface_tablegen_default
+#attr = #test.op_asm_attr_interface_tablegen_default<value = "test">
+
+func.func @test_op_asm_attr_interface() {
+  %1 = "test.result_name_from_type"() {attr = #attr} : () -> !test.op_asm_type_interface
+  return
+}

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

@@ -424,4 +424,15 @@ def TestOpAsmAttrInterfaceAttr : Test_Attr<"TestOpAsmAttrInterface",
   }];
 }
 
+// Test OpAsmAttrInterface from tablegen genMnemonicAlias option.
+def TestOpAsmAttrInterfaceTablegenDefaultAttr : Test_Attr<"TestOpAsmAttrInterfaceTablegenDefault"> {
+  let mnemonic = "op_asm_attr_interface_tablegen_default";
+  let parameters = (ins "mlir::StringAttr":$value);
+  let assemblyFormat = [{
+    `<` struct(params) `>`
+  }];
+
+  let genMnemonicAlias = 1;
+}
+
 #endif // TEST_ATTRDEFS

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

@@ -404,6 +404,11 @@ def TestTypeOpAsmTypeInterface : Test_Type<"TestTypeOpAsmTypeInterface",
   let mnemonic = "op_asm_type_interface";
 }
 
+def TestTypeOpAsmTypeInterfaceTablegenDefault : Test_Type<"TestTypeOpAsmTypeInterfaceTablegenDefault"> {
+  let mnemonic = "op_asm_type_interface_tablegen_default";
+  let genMnemonicAlias = 1;
+}
+
 def TestTensorType : Test_Type<"TestTensor",
     [Bufferization_TensorLikeTypeInterface, ShapedTypeInterface]> {
   let mnemonic = "test_tensor";

mlir/test/mlir-tblgen/attrdefs.td

@@ -172,3 +172,17 @@ def H_TestExtraClassAttr : TestAttr<"TestExtraClass"> {
 // DEF-LABEL: int TestExtraClassAttr::getFoo(int i) {
 // DEF: return i+1;
 // DEF-NEXT: }
+
+def I_TestGenMnemonicAliasAttr : TestAttr<"TestGenMnemonicAlias"> {
+  let mnemonic = "test_gen_mnemonic_alias";
+  let genMnemonicAlias = 1;
+}
+
+// DECL-LABEL: class TestGenMnemonicAliasAttr : public ::mlir::Attribute
+// DECL-SAME: ::mlir::OpAsmAttrInterface::Trait
+// DECL: ::mlir::OpAsmAliasResult getAlias(::llvm::raw_ostream &os) const;
+
+// DEF-LABEL: ::mlir::OpAsmAliasResult TestGenMnemonicAliasAttr::getAlias(::llvm::raw_ostream &os) const {
+// DEF-NEXT: os << "test_gen_mnemonic_alias";
+// DEF-NEXT: return ::mlir::OpAsmAliasResult::OverridableAlias;
+// DEF-NEXT: }

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

@@ -130,6 +130,12 @@ class DefGen {
   /// Emit a trait method.
   void emitTraitMethod(const InterfaceMethod &method);
 
+  //===--------------------------------------------------------------------===//
+  // OpAsm{Type,Attr}Interface Default Method Emission
+
+  /// Emit 'getAlias' method using mnemonic as alias.
+  void emitMnemonicAliasMethod();
+
   //===--------------------------------------------------------------------===//
   // Storage Class Emission
   void emitStorageClass();
@@ -215,6 +221,9 @@ DefGen::DefGen(const AttrOrTypeDef &def)
     emitAccessors();
   // Emit trait interface methods
   emitInterfaceMethods();
+  // Emit OpAsm{Type,Attr}Interface default methods
+  if (def.genMnemonicAlias())
+    emitMnemonicAliasMethod();
   defCls.finalize();
   // Emit a storage class if one is needed
   if (storageCls && def.genStorageClass())
@@ -229,11 +238,24 @@ void DefGen::createParentWithTraits() {
                                  ? strfmt("{0}::{1}", def.getStorageNamespace(),
                                           def.getStorageClassName())
                                  : strfmt("::mlir::{0}Storage", valueType));
-  for (auto &trait : def.getTraits()) {
-    defParent.addTemplateParam(
-        isa<NativeTrait>(&trait)
-            ? cast<NativeTrait>(&trait)->getFullyQualifiedTraitName()
-            : cast<InterfaceTrait>(&trait)->getFullyQualifiedTraitName());
+  SmallVector<std::string> traitNames =
+      llvm::to_vector(llvm::map_range(def.getTraits(), [](auto &trait) {
+        return isa<NativeTrait>(&trait)
+                   ? cast<NativeTrait>(&trait)->getFullyQualifiedTraitName()
+                   : cast<InterfaceTrait>(&trait)->getFullyQualifiedTraitName();
+      }));
+  llvm::for_each(traitNames, [&](auto &traitName) {
+    defParent.addTemplateParam(traitName);
+  });
+
+  // Add OpAsmInterface::Trait if we automatically generate mnemonic alias
+  // method.
+  std::string opAsmInterfaceTraitName =
+      strfmt("::mlir::OpAsm{0}Interface::Trait", defType);
+  if (def.genMnemonicAlias() && llvm::none_of(traitNames, [&](auto &traitName) {
+        return traitName == opAsmInterfaceTraitName;
+      })) {
+    defParent.addTemplateParam(opAsmInterfaceTraitName);
   }
   defCls.addParent(std::move(defParent));
 }
@@ -577,6 +599,22 @@ void DefGen::emitTraitMethod(const InterfaceMethod &method) {
                    std::move(params));
 }
 
+//===----------------------------------------------------------------------===//
+// OpAsm{Type,Attr}Interface Default Method Emission
+
+void DefGen::emitMnemonicAliasMethod() {
+  // If the mnemonic is not set, there is nothing to do.
+  if (!def.getMnemonic())
+    return;
+
+  // Emit the mnemonic alias method.
+  SmallVector<MethodParameter> params{{"::llvm::raw_ostream &", "os"}};
+  Method *m = defCls.addMethod<Method::Const>("::mlir::OpAsmAliasResult",
+                                              "getAlias", std::move(params));
+  m->body().indent() << strfmt("os << \"{0}\";\n", *def.getMnemonic())
+                     << "return ::mlir::OpAsmAliasResult::OverridableAlias;\n";
+}
+
 //===----------------------------------------------------------------------===//
 // Storage Class Emission
 //===----------------------------------------------------------------------===//