[mlir] Add OpAsmTypeInterface for pretty-print

Author: ZenithalHourlyRate

mlir/include/mlir/IR/CMakeLists.txt

@@ -1,7 +1,14 @@
-add_mlir_interface(OpAsmInterface)
 add_mlir_interface(SymbolInterfaces)
 add_mlir_interface(RegionKindInterface)
 
+set(LLVM_TARGET_DEFINITIONS OpAsmInterface.td)
+mlir_tablegen(OpAsmOpInterface.h.inc -gen-op-interface-decls)
+mlir_tablegen(OpAsmOpInterface.cpp.inc -gen-op-interface-defs)
+mlir_tablegen(OpAsmTypeInterface.h.inc -gen-type-interface-decls)
+mlir_tablegen(OpAsmTypeInterface.cpp.inc -gen-type-interface-defs)
+add_public_tablegen_target(MLIROpAsmInterfaceIncGen)
+add_dependencies(mlir-generic-headers MLIROpAsmInterfaceIncGen)
+
 set(LLVM_TARGET_DEFINITIONS BuiltinAttributes.td)
 mlir_tablegen(BuiltinAttributes.h.inc -gen-attrdef-decls)
 mlir_tablegen(BuiltinAttributes.cpp.inc -gen-attrdef-defs)

mlir/include/mlir/IR/OpAsmInterface.td

@@ -50,10 +50,27 @@ def OpAsmOpInterface : OpInterface<"OpAsmOpInterface"> {
         ```mlir
           %first_result, %middle_results:2, %0 = "my.op" ...
         ```
+
+        The default implementation uses `OpAsmTypeInterface` to get the name for
+        each result from its type.
+
+        If not all of the result types have `OpAsmTypeInterface`, the default implementation
+        does nothing, as the packing behavior should be decided by the operation itself.
       }],
       "void", "getAsmResultNames",
       (ins "::mlir::OpAsmSetValueNameFn":$setNameFn),
-      "", "return;"
+      "", [{
+        bool hasOpAsmTypeInterface = llvm::all_of($_op->getResults(), [&](Value result) {
+          return ::mlir::isa<::mlir::OpAsmTypeInterface>(result.getType());
+        });
+        if (!hasOpAsmTypeInterface)
+          return;
+        for (auto result : $_op->getResults()) {
+          auto setResultNameFn = [&](StringRef name) { setNameFn(result, name); };
+          auto opAsmTypeInterface = ::mlir::cast<::mlir::OpAsmTypeInterface>(result.getType());
+          opAsmTypeInterface.getAsmName(setResultNameFn);
+        }
+      }]
     >,
     InterfaceMethod<[{
         Get a special name to use when printing the block arguments for a region
@@ -64,7 +81,16 @@ def OpAsmOpInterface : OpInterface<"OpAsmOpInterface"> {
         "::mlir::Region&":$region,
         "::mlir::OpAsmSetValueNameFn":$setNameFn
       ),
-      "", "return;"
+      "", [{
+        for (auto &block : region) {
+          for (auto arg : block.getArguments()) {
+            if (auto opAsmTypeInterface = ::mlir::dyn_cast<::mlir::OpAsmTypeInterface>(arg.getType())) {
+              auto setArgNameFn = [&](StringRef name) { setNameFn(arg, name); };
+              opAsmTypeInterface.getAsmName(setArgNameFn);
+            }
+          }
+        }
+      }]
     >,
     InterfaceMethod<[{
         Get the name to use for a given block inside a region attached to this
@@ -109,6 +135,31 @@ def OpAsmOpInterface : OpInterface<"OpAsmOpInterface"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// OpAsmTypeInterface
+//===----------------------------------------------------------------------===//
+
+def OpAsmTypeInterface : TypeInterface<"OpAsmTypeInterface"> {
+  let description = [{
+    This interface provides hooks to interact with the AsmPrinter and AsmParser
+    classes.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Get a special name to use when printing value of this type.
+
+        For example, the default implementation of OpAsmOpInterface
+        will respect this method when printing the results of an operation
+        and/or block argument of it.
+      }],
+      "void", "getAsmName",
+      (ins "::mlir::OpAsmSetNameFn":$setNameFn), "", ";"
+    >,
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // ResourceHandleParameter
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/OpImplementation.h

@@ -734,7 +734,7 @@ class AsmParser {
   virtual OptionalParseResult parseOptionalInteger(APInt &result) = 0;
   virtual OptionalParseResult parseOptionalDecimalInteger(APInt &result) = 0;
 
- private:
+private:
   template <typename IntT, typename ParseFn>
   OptionalParseResult parseOptionalIntegerAndCheck(IntT &result,
                                                    ParseFn &&parseFn) {
@@ -756,7 +756,7 @@ class AsmParser {
     return success();
   }
 
- public:
+public:
   template <typename IntT>
   OptionalParseResult parseOptionalInteger(IntT &result) {
     return parseOptionalIntegerAndCheck(
@@ -1727,6 +1727,10 @@ class OpAsmParser : public AsmParser {
 // Dialect OpAsm interface.
 //===--------------------------------------------------------------------===//
 
+/// A functor used to set the name of the result. See 'getAsmResultNames' below
+/// for more details.
+using OpAsmSetNameFn = function_ref<void(StringRef)>;
+
 /// A functor used to set the name of the start of a result group of an
 /// operation. See 'getAsmResultNames' below for more details.
 using OpAsmSetValueNameFn = function_ref<void(Value, StringRef)>;
@@ -1820,7 +1824,9 @@ ParseResult parseDimensionList(OpAsmParser &parser,
 //===--------------------------------------------------------------------===//
 
 /// The OpAsmOpInterface, see OpAsmInterface.td for more details.
-#include "mlir/IR/OpAsmInterface.h.inc"
+#include "mlir/IR/OpAsmTypeInterface.h.inc"
+// put Attr/Type before Op
+#include "mlir/IR/OpAsmOpInterface.h.inc"
 
 namespace llvm {
 template <>

mlir/lib/IR/AsmPrinter.cpp

@@ -125,7 +125,8 @@ void OpAsmPrinter::printFunctionalType(Operation *op) {
 //===----------------------------------------------------------------------===//
 
 /// The OpAsmOpInterface, see OpAsmInterface.td for more details.
-#include "mlir/IR/OpAsmInterface.cpp.inc"
+#include "mlir/IR/OpAsmOpInterface.cpp.inc"
+#include "mlir/IR/OpAsmTypeInterface.cpp.inc"
 
 LogicalResult
 OpAsmDialectInterface::parseResource(AsmParsedResourceEntry &entry) const {

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

@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Test OpAsmOpInterface
+//===----------------------------------------------------------------------===//
+
+func.func @result_name_from_op_asm_type_interface() {
+  // CHECK-LABEL: @result_name_from_op_asm_type_interface
+  // CHECK: %op_asm_type_interface
+  %0 = "test.default_result_name"() : () -> !test.op_asm_type_interface
+  return
+}
+
+// -----
+
+func.func @result_name_pack_from_op_asm_type_interface() {
+  // CHECK-LABEL: @result_name_pack_from_op_asm_type_interface
+  // CHECK: %op_asm_type_interface{{.*}}, %op_asm_type_interface{{.*}}
+  // CHECK-NOT: :2
+  %0:2 = "test.default_result_name_packing"() : () -> (!test.op_asm_type_interface, !test.op_asm_type_interface)
+  return
+}
+
+// -----
+
+func.func @result_name_pack_do_nothing() {
+  // CHECK-LABEL: @result_name_pack_do_nothing
+  // CHECK: %0:2
+  %0:2 = "test.default_result_name_packing"() : () -> (i32, !test.op_asm_type_interface)
+  return
+}
+
+// -----
+
+func.func @block_argument_name_from_op_asm_type_interface() {
+  // CHECK-LABEL: @block_argument_name_from_op_asm_type_interface
+  // CHECK: ^bb0(%op_asm_type_interface
+  test.default_block_argument_name {
+    ^bb0(%arg0: !test.op_asm_type_interface):
+      "test.terminator"() : ()->()
+  }
+  return
+}

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

@@ -924,6 +924,30 @@ def CustomResultsNameOp
   let results = (outs Variadic<AnyInteger>:$r);
 }
 
+// This is used to test default implementation of OpAsmOpInterface::getAsmResultNames,
+// which uses OpAsmTypeInterface if available.
+def DefaultResultsNameOp
+ : TEST_Op<"default_result_name",
+           [OpAsmOpInterface]> {
+  let results = (outs AnyType:$r);
+}
+
+// This is used to test default implementation of OpAsmOpInterface::getAsmResultNames,
+// when there are multiple results, and not all of their type has OpAsmTypeInterface,
+// it should not set result name from OpAsmTypeInterface.
+def DefaultResultsNamePackingOp
+ : TEST_Op<"default_result_name_packing",
+           [OpAsmOpInterface]> {
+  let results = (outs AnyType:$r, AnyType:$s);
+}
+
+// This is used to test default implementation of OpAsmOpInterface::getAsmBlockArgumentNames,
+def DefaultBlockArgumentNameOp : TEST_Op<"default_block_argument_name",
+                                    [OpAsmOpInterface]> {
+  let regions = (region AnyRegion:$body);
+  let assemblyFormat = "regions attr-dict-with-keyword";
+}
+
 // This is used to test the OpAsmOpInterface::getDefaultDialect() feature:
 // operations nested in a region under this op will drop the "test." dialect
 // prefix.

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

@@ -398,4 +398,8 @@ def TestTypeVerification : Test_Type<"TestTypeVerification"> {
   let assemblyFormat = "`<` $param `>`";
 }
 
+def TestTypeOpAsmTypeInterface : Test_Type<"TestTypeOpAsmTypeInterface", [DeclareTypeInterfaceMethods<OpAsmTypeInterface, ["getAsmName"]>]> {
+  let mnemonic = "op_asm_type_interface";
+}
+
 #endif // TEST_TYPEDEFS

mlir/test/lib/Dialect/Test/TestTypes.cpp

@@ -531,3 +531,8 @@ void TestRecursiveAliasType::print(AsmPrinter &printer) const {
   }
   printer << ">";
 }
+
+void TestTypeOpAsmTypeInterfaceType::getAsmName(
+    OpAsmSetNameFn setNameFn) const {
+  setNameFn("op_asm_type_interface");
+}