[MLIR] Add option to print users of an operation as comment in the printer

[cpillmayer]

Summary

This allows printing the users of an operation as proposed in the git issue #53286.
To be able to refer to operations with no result, these operations are assigned an
ID in SSANameState.

[mehdi_amini]

Can you add tests to cover all this? I'd to see how it looks in practice!

Also please include a test with a user operation that customize its result names? An op like %q, %q_1, %q_2, %r = test.string_attr_pretty_name attributes {names = ["q", "q", "q", "r"]} (but that also accepts operands, you need to create one or modify this one to add an optional argument).

mlir/lib/IR/AsmPrinter.cpp

2720

Why wouldn't you print for an operation with operands and regions?

2725

Do you think explicit is better?

2728

What about grouping per results? If an op is like %a, %b = my.op, I could see useful to track it this way: %a, %b = my.op // users: (%42, %56), (%73, %45)

2729

Nit: Operation *
(we reserve auto to case when the type is obvious of the context, or it is more readable than spelling the type)

2732

You can check the return of the insert and save a lookup instead of calling contains separately

2746

(same as above)

[mehdi_amini]

This comment was removed by mehdi_amini.

[mehdi_amini]

Actually I patched your file to try it:

module {
  func @pretty_names(%arg0: i64) {
    // %arg0 // user: %q
    %c0_i64 = arith.constant 0 : i64 // user: %q
    %q, %q_0, %q_1, %r = "test.custom_result_name"(%c0_i64, %arg0) {names = ["q", "q", "q", "r"]} : (i64, i64) -> (i32, i32, i32, i32)
    return
  }
}

I think the block argument handling isn't ideal right now?

[mehdi_amini]

Slightly more elaborated:

$ ./bin/mlir-opt /tmp/test.mlir -mlir-print-operation-users
module {
  func @pretty_names(%arg0: i32, %arg1: i32) {
    // %arg0 // users: %q_2, %q
    // %arg1 // users: %q_6, %q_2
    %c0_i32 = arith.constant 0 : i32 // users: %q
    %q, %q_0, %q_1, %r = "test.custom_result_name"(%c0_i32, %c0_i32, %c0_i32, %arg0) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32) // users: %q_6, %q_2
    %q_2, %q_3, %q_4, %r_5 = "test.custom_result_name"(%q_0, %q, %arg1, %arg0) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32) // users: %q_6
    %q_6, %q_7, %q_8, %r_9 = "test.custom_result_name"(%q, %r_5, %q_2, %arg1) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32)
    return
  }
}

[cpillmayer]

Address part of comments

Print id for operation with regions and indicate operation with unused results.
Add suggested code improvements.

[cpillmayer]

mlir/lib/IR/AsmPrinter.cpp

2720

I thought of the operation too much in terms of a function. So I wanted to not print the id for a it as it would already have a name and the extra id comment would be redundant. But that doesn't make sense.

2725

Definitely makes sense whenever there is a result but no user. I added that.

2728

In this example it looks good. From my understanding the printer would always join the named values together and then we could get a result like this

%a:2 = my.op // users: (%1), (%1)
%1 = my.other(%a#0, %a#1) ...

I thought it is not great to have %1 in there twice.
How would you handle this?

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

2728

From my understanding the printer would always join the named values together

It does not always do it the op can control that.
Did you see the example I posted? It has this for example:

%q, %q_0, %q_1, %r = "test.custom_result_name"(%c0_i32, %c0_i32, %c0_i32, %arg0) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32) // users: %q_6, %q_2

But even if it collapsed, that is just cosmetic. I would instead wonder what is the best information to provide to the user when there are multiple results? That is: when there are multiple results would the important information be about the user of the operation or the users of an indvidual result? Expanding your example it could look like:

%a:2 = my.op // users: (%1 %2), (%1, %3)
%1 = my.other(%a#0, %a#1) ...
%2 = my.other(%a#0) ...
%3 = my.other(%a#1) ...

There aren't many ops with multiple results, scf.if may be one:

%if#2 = scf.if %cond -> (index, index) { // Users: (%2, %3, %4, %5) or // Users: (%2, %4), (%3, %4, %5)
  ...
  scf.yield %a, %b : index, index
} else {
  ...
  scf.yield %c, %d : index, index
}
%2 = foo(%if#0)
%3 = bar(%if#1)
%4 = foo(%if#0, %if#1)
%5 = bar(%if#1)

[cpillmayer]

In D124048#3460640, @mehdi_amini wrote:

Slightly more elaborated:

$ ./bin/mlir-opt /tmp/test.mlir -mlir-print-operation-users
module {
  func @pretty_names(%arg0: i32, %arg1: i32) {
    // %arg0 // users: %q_2, %q
    // %arg1 // users: %q_6, %q_2
    %c0_i32 = arith.constant 0 : i32 // users: %q
    %q, %q_0, %q_1, %r = "test.custom_result_name"(%c0_i32, %c0_i32, %c0_i32, %arg0) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32) // users: %q_6, %q_2
    %q_2, %q_3, %q_4, %r_5 = "test.custom_result_name"(%q_0, %q, %arg1, %arg0) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32) // users: %q_6
    %q_6, %q_7, %q_8, %r_9 = "test.custom_result_name"(%q, %r_5, %q_2, %arg1) {names = ["q", "q", "q", "r"]} : (i32, i32, i32, i32) -> (i32, i32, i32, i32)
    return
  }
}

Thanks for the comments @mehdi_amini. How would you test the behaviour with the custom result names? For me, mlir-opt does not behave like you show in your comment and always gives anonymous names to the results.

[mehdi_amini]

I patched MLIR for this:

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 5cfefbb55971..699a365a1e2f 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -1168,6 +1168,15 @@ void StringAttrPrettyNameOp::getAsmResultNames(
         setNameFn(getResult(i), str.getValue());
 }
 
+void CustomResultsNameOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  auto value = getNames();
+  for (size_t i = 0, e = value.size(); i != e; ++i)
+    if (auto str = value[i].dyn_cast<StringAttr>())
+      if (!str.getValue().empty())
+        setNameFn(getResult(i), str.getValue());
+}
+
 //===----------------------------------------------------------------------===//
 // ResultTypeWithTraitOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index f42fb8edcc18..c2aa2cd49fa0 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -727,11 +727,27 @@ def OIListAllowedLiteral : TEST_Op<"oilist_allowed_literal"> {
 def StringAttrPrettyNameOp
  : TEST_Op<"string_attr_pretty_name",
            [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
-  let arguments = (ins StrArrayAttr:$names);
+  let arguments = (ins 
+    Optional<I64>:$optional,
+    StrArrayAttr:$names
+  );
   let results = (outs Variadic<I32>:$r);
   let hasCustomAssemblyFormat = 1;
 }
 
+
+// This is used to test encoding of a string attribute into an SSA name of a
+// pretty printed value name.
+def CustomResultsNameOp
+ : TEST_Op<"custom_result_name",
+           [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+  let arguments = (ins 
+    Variadic<AnyInteger>:$optional,
+    StrArrayAttr:$names
+  );
+  let results = (outs Variadic<AnyInteger>:$r);
+}
+
 // This is used to test the OpAsmOpInterface::getDefaultDialect() feature:
 // operations nested in a region under this op will drop the "test." dialect
 // prefix.

[cpillmayer]

mlir/lib/IR/AsmPrinter.cpp

2728

You are right, thanks for the example. I think the user would prefer to check where a result is used instead of where any result from an operation is used.

[cpillmayer]

Address comments

• Add tests
• Add grouping of users based on results
• Improve printing of BlockArgument users

[cpillmayer]

In D124048#3462780, @mehdi_amini wrote:

I patched MLIR for this:
...

Thanks a lot for your help. This made it much clearer for me. I added your operation to create the tests.

[Mogball]

mlir/include/mlir/IR/OperationSupport.h

780

Add a function so that the flag can be set programmatically?

801

It's not just the users of operations but block arguments too. The comment and the names of everything should reflect that. printValueUsers?

mlir/lib/IR/AsmPrinter.cpp

1020

In what scenarios would this case be triggered?

2599

2602

2727

Fit this comment to 80 characters long lines.

2732

Drop the explicit static size specifier unless you have a good reason to have it

2733

spell out the auto

2775

2776

2792

2882

spell out the auto

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

62 ↗

(On Diff #424252)

nit

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

1173

spell out the auto

1174

prefer unsigned

[mehdi_amini]

LGTM overall!

mlir/include/mlir/IR/OperationSupport.h

780

+1: looks at the setters style line 736-753

mlir/lib/IR/AsmPrinter.cpp

1020

Invalid IR?

2732

That work only with SmallVector AFAIK

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

1174

prefer int! :)

for (int i : llvm::seq<int>(0, value.size())

Or:

for (auto nameAndResult : llvm::zip(getNames(), getResults()))
  if (auto str = std::get<0>(nameAndResult).dyn_cast<StringAttr>())
      if (!str.getValue().empty())
        setNameFn(std::get<1>(nameAndResult), str.getValue());

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1158

Can you reflow this comment?

2718

Please drop comments from function definitions, prefer having them only on the declaration.

2720

Please spell out auto here.

2724

Prefer op->use_empty() over op->getUsers().empty()

2727

Can you please fix the flow of the comments? clang-format can auto format the comments so that they wrap properly.

2727

Should print feels a bit weird to read, can you tweak this comment?

2729

Please use punctuation for documentation.

2733

Please wrap the for in braces as well.

2739

Please add punctuation to the end of comments.

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

1174

(I would say size_t is fine and consistent with the codebase)

[cpillmayer]

mlir/lib/IR/AsmPrinter.cpp

1020

This case should never be reached right now imo. An operation id is only printed if the operation has no result, but in that case an id should have been assigned. I added the output in this case to be safe and so that this method always has reasonable output.

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

1174

This method is the same as for StringAttrPrettyNameOp::getAsmResultNames. Maybe it makes sense to change both at once in a follow up?

[cpillmayer]

Address comments

[mehdi_amini]

LGTM!
Do you have commit access?

[cpillmayer]

In D124048#3468089, @mehdi_amini wrote:

LGTM!
Do you have commit access?

Great, thanks again for the help and the feedback to all of you!
I don't have commit access, so it would be great if you could land the patch for me. Please use "Christoph Pillmayer [e-mail]" for the commit.
And also credit yourself for the new test op you provided me with, if that's possible?

[mehdi_amini]

In D124048#3468331, @cpillmayer wrote:

I don't have commit access, so it would be great if you could land the patch for me. Please use "Christoph Pillmayer [e-mail]" for the commit.

Sorry, I applied the patch with arc and forgot to check that the name was right, it appears as Author: cpillmayer <[e-mail]>, hopefully this is good enough?

[cpillmayer]

In D124048#3469663, @mehdi_amini wrote:

In D124048#3468331, @cpillmayer wrote:

I don't have commit access, so it would be great if you could land the patch for me. Please use "Christoph Pillmayer [e-mail]" for the commit.

Sorry, I applied the patch with arc and forgot to check that the name was right, it appears as Author: cpillmayer <[e-mail]>, hopefully this is good enough?

Yes, that's good as well. No worries, Thanks.

mlir/include/mlir/IR/OperationSupport.h

public:
   OpPrintingFlags &assumeVerified();
 
   /// Use local scope when printing the operation. This allows for using the
   /// printer in a more localized and thread-safe setting, but may not
   /// necessarily be identical to what the IR will look like when dumping
   /// the full module.
   OpPrintingFlags &useLocalScope();
 
+  /// Print users of values as comments.
+  OpPrintingFlags &printValueUsers();
+
   /// Return if the given ElementsAttr should be elided.
   bool shouldElideElementsAttr(ElementsAttr attr) const;
 
   /// Return the size limit for printing large ElementsAttr.
   Optional<int64_t> getLargeElementsAttrLimit() const;
 
   /// Return if debug information should be printed.
   bool shouldPrintDebugInfo() const;
 
   /// Return if debug information should be printed in the pretty form.
   bool shouldPrintDebugInfoPrettyForm() const;
 
   /// Return if operations should be printed in the generic form.
   bool shouldPrintGenericOpForm() const;
 
   /// Return if operation verification should be skipped.
   bool shouldAssumeVerified() const;
 
   /// Return if the printer should use local scope when dumping the IR.
   bool shouldUseLocalScope() const;
 
+  /// Return if the printer should print users of values.
+  bool shouldPrintValueUsers() const;
+
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// the upper limit.
   Optional<int64_t> elementsAttrElementLimit;
 
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;
 
   /// Print operations in the generic form.
   bool printGenericOpFormFlag : 1;
 
   /// Skip operation verification.
   bool assumeVerifiedFlag : 1;
 
   /// Print operations with numberings local to the current operation.
   bool printLocalScope : 1;
+
+  /// Print users of values.
+  bool printValueUsersFlag : 1;
 };
 
 //===----------------------------------------------------------------------===//
 // Operation Value-Iterators
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // OperandRange

mlir/lib/IR/AsmPrinter.cpp

llvm::cl::opt<bool> assumeVerifiedOpt{
       "mlir-print-assume-verified", llvm::cl::init(false),
       llvm::cl::desc("Skip op verification when using custom printers"),
       llvm::cl::Hidden};
 
   llvm::cl::opt<bool> printLocalScopeOpt{
       "mlir-print-local-scope", llvm::cl::init(false),
       llvm::cl::desc("Print with local scope and inline information (eliding "
                      "aliases for attributes, types, and locations")};
+
+  llvm::cl::opt<bool> printValueUsers{
+      "mlir-print-value-users", llvm::cl::init(false),
+      llvm::cl::desc(
+          "Print users of operation results and block arguments as a comment")};
 };
 } // namespace
 
 static llvm::ManagedStatic<AsmPrinterOptions> clOptions;
 
 /// Register a set of useful command-line options that can be used to configure
 /// various flags within the AsmPrinter.
 void mlir::registerAsmPrinterCLOptions() {
   // Make sure that the options struct has been initialized.
   *clOptions;
 }
 
 /// Initialize the printing flags with default supplied by the cl::opts above.
 OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
       printGenericOpFormFlag(false), assumeVerifiedFlag(false),
-      printLocalScope(false) {
+      printLocalScope(false), printValueUsersFlag(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;
   assumeVerifiedFlag = clOptions->assumeVerifiedOpt;
   printLocalScope = clOptions->printLocalScopeOpt;
+  printValueUsersFlag = clOptions->printValueUsers;
 }
 
 /// Enable the elision of large elements attributes, by printing a '...'
 /// instead of the element data, when the number of elements is greater than
 /// `largeElementLimit`. Note: The IR generated with this option is not
 /// parsable.
 OpPrintingFlags &
 OpPrintingFlags::elideLargeElementsAttrs(int64_t largeElementLimit) {

 /// Use local scope when printing the operation. This allows for using the
 /// printer in a more localized and thread-safe setting, but may not necessarily
 /// be identical of what the IR will look like when dumping the full module.
 OpPrintingFlags &OpPrintingFlags::useLocalScope() {
   printLocalScope = true;
   return *this;
 }
 
+/// Print users of values as comments.
+OpPrintingFlags &OpPrintingFlags::printValueUsers() {
+  printValueUsersFlag = true;
+  return *this;
+}
+
 /// Return if the given ElementsAttr should be elided.
 bool OpPrintingFlags::shouldElideElementsAttr(ElementsAttr attr) const {
   return elementsAttrElementLimit.hasValue() &&
          *elementsAttrElementLimit < int64_t(attr.getNumElements()) &&
          !attr.isa<SplatElementsAttr>();
 }
 
 /// Return the size limit for printing large ElementsAttr.

 /// Return if operation verification should be skipped.
 bool OpPrintingFlags::shouldAssumeVerified() const {
   return assumeVerifiedFlag;
 }
 
 /// Return if the printer should use local scope when dumping the IR.
 bool OpPrintingFlags::shouldUseLocalScope() const { return printLocalScope; }
 
+/// Return if the printer should print users of values.
+bool OpPrintingFlags::shouldPrintValueUsers() const {
+  return printValueUsersFlag;
+}
+
 /// Returns true if an ElementsAttr with the given number of elements should be
 /// printed with hex.
 static bool shouldPrintElementsAttrWithHex(int64_t numElements) {
   // Check to see if a command line option was provided for the limit.
   if (clOptions.isConstructed()) {
     if (clOptions->printElementsAttrWithHexIfLarger.getNumOccurrences()) {
       // -1 is used to disable hex printing.
       if (clOptions->printElementsAttrWithHexIfLarger == -1)
public:
 
   SSANameState(Operation *op, const OpPrintingFlags &printerFlags);
 
   /// Print the SSA identifier for the given value to 'stream'. If
   /// 'printResultNo' is true, it also presents the result number ('#' number)
   /// of this value.
   void printValueID(Value value, bool printResultNo, raw_ostream &stream) const;
 
+  /// Print the operation identifier.
+  void printOperationID(Operation *op, raw_ostream &stream) const;
+
   /// Return the result indices for each of the result groups registered by this
   /// operation, or empty if none exist.
   ArrayRef<int> getOpResultGroups(Operation *op);
 
   /// Get the info for the given block.
   BlockInfo getBlockInfo(Block *block);
 
   /// Renumber the arguments for the specified region to the same names as the
private:
   /// conflicts, it is automatically renamed.
   StringRef uniqueValueName(StringRef name);
 
   /// This is the value ID for each SSA value. If this returns NameSentinel,
   /// then the valueID has an entry in valueNames.
   DenseMap<Value, unsigned> valueIDs;
   DenseMap<Value, StringRef> valueNames;
 
+  /// When printing users of values, an operation without a result might
+  /// be the user. This map holds ids for such operations.
+  DenseMap<Operation *, unsigned> operationIDs;
+
   /// This is a map of operations that contain multiple named result groups,
   /// i.e. there may be multiple names for the results of the operation. The
   /// value of this map are the result numbers that start a result group.
   DenseMap<Operation *, SmallVector<int, 1>> opResultGroups;
 
   /// This maps blocks to there visitation number in the current region as well
   /// as the string representing their name.
   DenseMap<Block *, BlockInfo> blockNames;
if (it->second != NameSentinel) {
     assert(nameIt != valueNames.end() && "Didn't have a name entry?");
     stream << nameIt->second;
   }
 
   if (resultNo.hasValue() && printResultNo)
     stream << '#' << resultNo;
 }
 
+void SSANameState::printOperationID(Operation *op, raw_ostream &stream) const {
+  auto it = operationIDs.find(op);
+  if (it == operationIDs.end()) {
+    stream << "<<UNKOWN OPERATION>>";
+  } else {
+    stream << '%' << it->second;
+  }
+}
+
 ArrayRef<int> SSANameState::getOpResultGroups(Operation *op) {
   auto it = opResultGroups.find(op);
   return it == opResultGroups.end() ? ArrayRef<int>() : it->second;
 }
 
 BlockInfo SSANameState::getBlockInfo(Block *block) {
   auto it = blockNames.find(block);
   BlockInfo invalidBlock{-1, "INVALIDBLOCK"};
void SSANameState::numberValuesInOp(Operation &op) {
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (OpAsmOpInterface asmInterface = dyn_cast<OpAsmOpInterface>(&op)) {
       asmInterface.getAsmBlockNames(setBlockNameFn);
       asmInterface.getAsmResultNames(setResultNameFn);
     }
   }
 
   unsigned numResults = op.getNumResults();
-  if (numResults == 0)
+  if (numResults == 0) {
+    // If value users should be printed, operations with no result need an id.
+    if (printerFlags.shouldPrintValueUsers()) {
+      if (operationIDs.try_emplace(&op, nextValueID).second)
+        ++nextValueID;
+    }
     return;
+  }
   Value resultBegin = op.getResult(0);
 
   // If the first result wasn't numbered, give it a default number.
   if (valueIDs.try_emplace(resultBegin, nextValueID).second)
     ++nextValueID;
 
   // If this operation has multiple result groups, mark it.
   if (resultGroups.size() != 1) {
public:
   /// operation of the block is not printed.
   void print(Block *block, bool printBlockArgs = true,
              bool printBlockTerminator = true);
 
   /// Print the ID of the given value, optionally with its result number.
   void printValueID(Value value, bool printResultNo = true,
                     raw_ostream *streamOverride = nullptr) const;
 
+  /// Print the ID of the given operation.
+  void printOperationID(Operation *op,
+                        raw_ostream *streamOverride = nullptr) const;
+
   //===--------------------------------------------------------------------===//
   // OpAsmPrinter methods
   //===--------------------------------------------------------------------===//
 
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   void printNewline() override {
     os << newLine;
public:
   void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                               ValueRange operands) override;
 
   /// Print the given affine expression with the symbol and dimension operands
   /// printed inline with the expression.
   void printAffineExprOfSSAIds(AffineExpr expr, ValueRange dimOperands,
                                ValueRange symOperands) override;
 
+  /// Print users of this operation or id of this operation if it has no result.
+  void printUsersComment(Operation *op);
+
+  /// Print users of this block arg.
+  void printUsersComment(BlockArgument arg);
+
+  /// Print the users of a value.
+  void printValueUsers(Value value);
+
+  /// Print either the ids of the result values or the id of the operation if
+  /// the operation has no results.
+  void printUserIDs(Operation *user, bool prefixComma = false);
+
 private:
   // Contains the stack of default dialects to use when printing regions.
   // A new dialect is pushed to the stack before parsing regions nested under an
   // operation implementing `OpAsmOpInterface`, and popped when done. At the
   // top-level we start with "builtin" as the default, so that the top-level
   // `module` operation prints as-is.
   SmallVector<StringRef> defaultDialectStack{"builtin"};
 

 
 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 (printerFlags.shouldPrintValueUsers())
+    printUsersComment(op);
 }
 
 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;
if (Dialect *dialect = op->getDialect()) {
       }
     }
   }
 
   // Otherwise print with the generic assembly form.
   printGenericOp(op, /*printOpName=*/true);
 }
 
+void OperationPrinter::printUsersComment(Operation *op) {
+  unsigned numResults = op->getNumResults();
+  if (!numResults && op->getNumOperands()) {
+    os << " // id: ";
+    printOperationID(op);
+  } else if (numResults && op->use_empty()) {
+    os << " // unused";
+  } else if (numResults && !op->use_empty()) {
+    // Print "user" if the operation has one result used to compute one other
+    // result, or is used in one operation with no result.
+    unsigned usedInNResults = 0;
+    unsigned usedInNOperations = 0;
+    SmallPtrSet<Operation *, 1> userSet;
+    for (Operation *user : op->getUsers()) {
+      if (userSet.insert(user).second) {
+        ++usedInNOperations;
+        usedInNResults += user->getNumResults();
+      }
+    }
+
+    // We already know that users is not empty.
+    bool exactlyOneUniqueUse =
+        usedInNResults <= 1 && usedInNOperations <= 1 && numResults == 1;
+    os << " // " << (exactlyOneUniqueUse ? "user" : "users") << ": ";
+    bool shouldPrintBrackets = numResults > 1;
+    auto printOpResult = [&](OpResult opResult) {
+      if (shouldPrintBrackets)
+        os << "(";
+      printValueUsers(opResult);
+      if (shouldPrintBrackets)
+        os << ")";
+    };
+
+    interleaveComma(op->getResults(), printOpResult);
+  }
+}
+
+void OperationPrinter::printUsersComment(BlockArgument arg) {
+  os << "// ";
+  printValueID(arg);
+  if (arg.use_empty()) {
+    os << " is unused";
+  } else {
+    os << " is used by ";
+    printValueUsers(arg);
+  }
+  os << newLine;
+}
+
+void OperationPrinter::printValueUsers(Value value) {
+  if (value.use_empty())
+    os << "unused";
+
+  // One value might be used as the operand of an operation more than once.
+  // Only print the operations results once in that case.
+  SmallPtrSet<Operation *, 1> userSet;
+  for (auto &indexedUser : enumerate(value.getUsers())) {
+    if (userSet.insert(indexedUser.value()).second)
+      printUserIDs(indexedUser.value(), indexedUser.index());
+  }
+}
+
+void OperationPrinter::printUserIDs(Operation *user, bool prefixComma) {
+  if (prefixComma)
+    os << ", ";
+
+  if (!user->getNumResults()) {
+    printOperationID(user);
+  } else {
+    interleaveComma(user->getResults(),
+                    [this](Value result) { printValueID(result); });
+  }
+}
+
 void OperationPrinter::printGenericOp(Operation *op, bool printOpName) {
   if (printOpName) {
     os << '"';
     printEscapedString(op->getName().getStringRef(), os);
     os << '"';
   }
   os << '(';
   interleaveComma(op->getOperands(), [&](Value value) { printValueID(value); });
   os << ')';
if (!block->getParent()) {
       os << "  // " << predIDs.size() << " preds: ";
 
       interleaveComma(predIDs, [&](BlockInfo pred) { os << pred.name; });
     }
     os << newLine;
   }
 
   currentIndent += indentWidth;
+
+  if (printerFlags.shouldPrintValueUsers()) {
+    for (BlockArgument arg : block->getArguments()) {
+      os.indent(currentIndent);
+      printUsersComment(arg);
+    }
+  }
+
   bool hasTerminator =
       !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   auto range = llvm::make_range(
       block->begin(),
       std::prev(block->end(),
                 (!hasTerminator || printBlockTerminator) ? 0 : 1));
   for (auto &op : range) {
     print(&op);
     os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
 void OperationPrinter::printValueID(Value value, bool printResultNo,
                                     raw_ostream *streamOverride) const {
   state->getSSANameState().printValueID(value, printResultNo,
                                         streamOverride ? *streamOverride : os);
 }
 
+void OperationPrinter::printOperationID(Operation *op,
+                                        raw_ostream *streamOverride) const {
+  state->getSSANameState().printOperationID(op, streamOverride ? *streamOverride
+                                                               : os);
+}
+
 void OperationPrinter::printSuccessor(Block *successor) {
   printBlockName(successor);
 }
 
 void OperationPrinter::printSuccessorAndUseList(Block *successor,
                                                 ValueRange succOperands) {
   printBlockName(successor);
   if (succOperands.empty())

mlir/test/IR/print-value-users.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-value-users -split-input-file %s | FileCheck %s
+
+module {
+    // CHECK: %[[ARG0:.+]]: i32, %[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32
+    func @foo(%arg0: i32, %arg1: i32, %arg3: i32) -> i32 {
+        // CHECK-NEXT: // %[[ARG0]] is used by %[[ARG0U1:.+]], %[[ARG0U2:.+]], %[[ARG0U3:.+]]
+        // CHECK-NEXT: // %[[ARG1]] is used by %[[ARG1U1:.+]], %[[ARG1U2:.+]]
+        // CHECK-NEXT: // %[[ARG2]] is unused
+        // CHECK-NEXT: test.noop
+        // CHECK-NOT: // unused
+        "test.noop"() : () -> ()
+        // When no result is produced, an id should be printed.
+        // CHECK-NEXT: // id: %[[ARG0U3]]
+        "test.no_result"(%arg0) {} : (i32) -> ()
+        // Check for unused result.
+        // CHECK-NEXT: %[[ARG0U2]] = 
+        // CHECK-SAME: // unused
+        %1 = "test.unused_result"(%arg0, %arg1) {} : (i32, i32) -> i32
+        // Check that both users are printed.
+        // CHECK-NEXT: %[[ARG0U1]] = 
+        // CHECK-SAME: // users: %[[A:.+]]#0, %[[A]]#1
+        %2 = "test.one_result"(%arg0, %arg1) {} : (i32, i32) -> i32
+        // For multiple results, users should be grouped per result.
+        // CHECK-NEXT: %[[A]]:2 = 
+        // CHECK-SAME: // users: (%[[B:.+]], %[[C:.+]]), (%[[B]], %[[D:.+]])
+        %3:2 = "test.many_results"(%2) {} : (i32) -> (i32, i32)
+        // Two results are produced, but there is only one user.
+        // CHECK-NEXT: // users:
+        %7:2 = "test.many_results"() : () -> (i32, i32)
+        // CHECK-NEXT: %[[C]] =
+        // Result is used twice in next operation but it produces only one result.
+        // CHECK-SAME: // user:
+        %4 = "test.foo"(%3#0) {} : (i32) -> i32
+        // CHECK-NEXT: %[[D]] =
+        %5 = "test.foo"(%3#1, %4, %4) {} : (i32, i32, i32) -> i32
+        // CHECK-NEXT: %[[B]] =
+        // Result is not used in any other result but in two operations.
+        // CHECK-SAME: // users:
+        %6 = "test.foo"(%3#0, %3#1) {} : (i32, i32) -> i32
+        "test.no_result"(%6) {} : (i32) -> ()
+        "test.no_result"(%7#0) : (i32) -> ()
+        return %6: i32
+    }
+}
+
+// -----
+
+module {
+    // Check with nested operation.
+    // CHECK: %[[CONSTNAME:.+]] = arith.constant
+    %0 = arith.constant 42 : i32
+    %test = "test.outerop"(%0) ({
+        // CHECK: "test.innerop"(%[[CONSTNAME]]) : (i32) -> () // id: %
+        "test.innerop"(%0) : (i32) -> ()
+    // CHECK: (i32) -> i32 // users: %r, %s, %p, %p_0, %q
+    }): (i32) -> i32
+
+    // Check named results.
+    // CHECK-NEXT: // users: (%u, %v), (unused), (%u, %v, %r, %s)
+    %p:2, %q = "test.custom_result_name"(%test) {names = ["p", "p", "q"]} : (i32) -> (i32, i32, i32)
+    // CHECK-NEXT: // users: (unused), (%u, %v)
+    %r, %s = "test.custom_result_name"(%q#0, %q#0, %test) {names = ["r", "s"]} : (i32, i32, i32) -> (i32, i32)
+    // CHECK-NEXT: // unused
+    %u, %v = "test.custom_result_name"(%s, %q#0, %p) {names = ["u", "v"]} : (i32, i32, i32) -> (i32, i32)
+}

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

void StringAttrPrettyNameOp::getAsmResultNames(
 
   auto value = getNames();
   for (size_t i = 0, e = value.size(); i != e; ++i)
     if (auto str = value[i].dyn_cast<StringAttr>())
       if (!str.getValue().empty())
         setNameFn(getResult(i), str.getValue());
 }
 
+void CustomResultsNameOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  ArrayAttr value = getNames();
+  for (size_t i = 0, e = value.size(); i != e; ++i)
+    if (auto str = value[i].dyn_cast<StringAttr>())
+      if (!str.getValue().empty())
+        setNameFn(getResult(i), str.getValue());
+}
+
 //===----------------------------------------------------------------------===//
 // ResultTypeWithTraitOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ResultTypeWithTraitOp::verify() {
   if ((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
     return success();
   return emitError("result type should have trait 'TestTypeTrait'");

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

 def StringAttrPrettyNameOp
  : TEST_Op<"string_attr_pretty_name",
            [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
   let arguments = (ins StrArrayAttr:$names);
   let results = (outs Variadic<I32>:$r);
   let hasCustomAssemblyFormat = 1;
 }
 
+
+// This is used to test encoding of a string attribute into an SSA name of a
+// pretty printed value name.
+def CustomResultsNameOp
+ : TEST_Op<"custom_result_name",
+           [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+  let arguments = (ins 
+    Variadic<AnyInteger>:$optional,
+    StrArrayAttr:$names
+  );
+  let results = (outs Variadic<AnyInteger>:$r);
+}
+
 // This is used to test the OpAsmOpInterface::getDefaultDialect() feature:
 // operations nested in a region under this op will drop the "test." dialect
 // prefix.
 def DefaultDialectOp : TEST_Op<"default_dialect", [OpAsmOpInterface]> {
  let regions = (region AnyRegion:$body);
   let extraClassDeclaration = [{
     static ::llvm::StringRef getDefaultDialect() {
       return "test";