Add a `printTrailingComment` method on OpAsmOpInterface

[mehdi_amini]

Summary

This is a demo for: https://llvm.discourse.group/t/adding-comments-to-declarative-assembly-format/5906/

It isn't fleshed out and requires more thoughts: again this is just a demo to help the discussion on Discourse.

[stephenneuendorffer]

Thanks!

mlir/lib/IR/AsmPrinter.cpp

2532–2535

braces not needed

[stephenneuendorffer]

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

666

This op prints....

[rriddle]

I'd rather not have trailing in the name. If we have multiple things adding comments, how is this going to interact?

[mehdi_amini]

In D117942#3263160, @rriddle wrote:

I'd rather not have trailing in the name. If we have multiple things adding comments, how is this going to interact?

I used trailing to express that it comes after the op completes (after trailing locations), and also because nothing can come after a trailing comment on the same line.
We don't support /*.. */ comments anyway right?
I don't think anything else can add comments today, but if we added some in the future it shouldn't be an issue, you can accumulate trailing comments (// comment1 // comment2) :)

[rriddle]

In D117942#3263163, @mehdi_amini wrote:

In D117942#3263160, @rriddle wrote:

I'd rather not have trailing in the name. If we have multiple things adding comments, how is this going to interact?

you can accumulate trailing comments (// comment1 // comment2) :)

(assuming this is multiple comments on 1 line) -1 there, that looks clunky (especially if the comments are more than a few characters) and would likely be very unreadable.

That being said, I think right now your documentation leaves it open ended what "after the operation" means; i.e. I guess it could mean on separate line/same line/etc.
I would prefer we don't guarantee that the comment is on the same line as the operation, maybe we try to do that when we can (maybe one line and one comment?)?

mlir/include/mlir/IR/OpAsmInterface.td

65–71

Should this just be limited to the operation interface? Why not also add to the dialect interface? In some sense, I expect we may also end up allowing injection of comments from users outside of interfaces (though not important now).

66

Can you expand this a bit? Giving more detail, examples, etc?

It also isn't clear here if the comment is allowed to span multiple lines (which I guess not right now from the implementation?).

mlir/lib/IR/AsmPrinter.cpp

2532–2534

Should we always print comments?

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

65–71

When would the dialect interface be involved? If it is on a per-op basis, then the dialect can always implement the fallback for the op right?

mlir/lib/IR/AsmPrinter.cpp

2532–2534

You mean like: not printing them in the generic format for example?

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

65–71

Ping on this River?

[mehdi_amini]

Discussed this with River:

• We want something more general: inject information to be printed as comment before, after, or trailing the operation.
• The interface is restricting us to a single information to be printed per operation. We can have this "print user count" added as an external model and also add another one. We likely want something that composes better.

mlir/include/mlir/IR/OpAsmInterface.td

InterfaceMethod<[{
       }],
       "void", "getAsmBlockArgumentNames",
       (ins
         "::mlir::Region&":$region,
         "::mlir::OpAsmSetValueNameFn":$setNameFn
       ),
       "", "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
       prefix from the operation name, for example it would be possible to omit
       the `spv.` prefix from all operations within a SpirV module if this method
       returned `spv`. The default implementation returns an empty string which
       is ignored.
       }],
       "StringRef", "getDefaultDialect", (ins), "", "return \"\";"
     >,
   ];
 }
 
 #endif // MLIR_OPASMINTERFACE

mlir/lib/IR/AsmPrinter.cpp

 
 void OperationPrinter::print(Operation *op) {
   // Track the location of this operation.
   state->registerOperationLocation(op, newLine.curLine, currentIndent);
 
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
+  if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op)) {
+    asmInterface.printTrailingComment(os);
+  }
 }
 
 void OperationPrinter::printOperation(Operation *op) {
   if (size_t numResults = op->getNumResults()) {
     auto printResultGroup = [&](size_t resultNo, size_t resultCount) {
       printValueID(op->getResult(resultNo), /*printResultNo=*/false);
       if (resultCount > 1)
         os << ':' << resultCount;
     };

mlir/test/IR/print-op-generic.mlir

 // # 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
 }
+
+

mlir/test/lib/Dialect/Test/TestDialect.h

 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #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;
 class RewritePatternSet;
 } // namespace mlir
 
 #include "TestOpInterfaces.h.inc"
 #include "TestOpStructs.h.inc"

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

let extraClassDeclaration = [{
     static ::llvm::StringRef getDefaultDialect() {
       return "test";
     }
     void getAsmResultNames(::llvm::function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) {}
   }];
   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);
 
   let verifier = [{
     if((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
       return success();
     return this->emitError("result type should have trait 'TestTypeTrait'");