diff --git a/mlir/include/mlir/IR/OpAsmInterface.td b/mlir/include/mlir/IR/OpAsmInterface.td --- a/mlir/include/mlir/IR/OpAsmInterface.td +++ b/mlir/include/mlir/IR/OpAsmInterface.td @@ -62,6 +62,13 @@ ), "", "return;" >, + InterfaceMethod<[{ + Invoked after an op is printed, allowing to print a comment. + }], + "void", "printTrailingComment", + (ins "::llvm::raw_ostream &":$os), + "", "return;" + >, StaticInterfaceMethod<[{ Return the default dialect used when printing/parsing operations in regions nested under this operation. This allows for eliding the dialect diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -2529,6 +2529,9 @@ os.indent(currentIndent); printOperation(op); printTrailingLocation(op->getLoc()); + if (auto asmInterface = dyn_cast(op)) { + asmInterface.printTrailingComment(os); + } } void OperationPrinter::printOperation(Operation *op) { diff --git a/mlir/test/IR/print-op-generic.mlir b/mlir/test/IR/print-op-generic.mlir --- a/mlir/test/IR/print-op-generic.mlir +++ b/mlir/test/IR/print-op-generic.mlir @@ -1,13 +1,16 @@ // # RUN: mlir-opt %s | FileCheck %s // # RUN: mlir-opt %s --mlir-print-op-generic | FileCheck %s --check-prefix=GENERIC -// CHECK-LABEL: func @pretty_names + +// CHECK-LABEL: func @print_comment // CHECK-GENERIC: "func"() -func @pretty_names() { - %x = test.string_attr_pretty_name - // CHECK: %x = test.string_attr_pretty_name - // GENERIC: %0 = "test.string_attr_pretty_name"() - return - // CHECK: return - // GENERIC: "std.return"() +func @print_comment() -> (i32, i32, i32) { + // CHECK: %0 = test.comment_user_count + // CHECK-SAME: // 3 uses + %0 = test.comment_user_count + // GENERIC: %0 = "test.comment_user_count"() + // GENERIC-SAME: // 3 uses + return %0, %0, %0 : i32, i32, i32 } + + diff --git a/mlir/test/lib/Dialect/Test/TestDialect.h b/mlir/test/lib/Dialect/Test/TestDialect.h --- a/mlir/test/lib/Dialect/Test/TestDialect.h +++ b/mlir/test/lib/Dialect/Test/TestDialect.h @@ -33,6 +33,7 @@ #include "mlir/Interfaces/DerivedAttributeOpInterface.h" #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "llvm/Support/raw_ostream.h" namespace mlir { class DLTIDialect; diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -663,6 +663,20 @@ let assemblyFormat = "regions attr-dict-with-keyword"; } +// This op print the count of the number of uses as a comment after it is +// printed, even in the generic print. +def CommentUserCountOp : TEST_Op<"comment_user_count", [OpAsmOpInterface]> { + let results = (outs I32:$r); + let extraClassDeclaration = [{ + void printTrailingComment(llvm::raw_ostream &os) { + auto uses = getR().getUses(); + int count = std::distance(uses.begin(), uses.end()); + os << " // " << count << " uses"; + } + }]; + let assemblyFormat = "attr-dict-with-keyword"; +} + // This operation requires its return type to have the trait 'TestTypeTrait'. def ResultTypeWithTraitOp : TEST_Op<"result_type_with_trait", []> { let results = (outs AnyType);