[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

2711

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

2716

Do you think explicit is better?

2719

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)

2720

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

2723

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

2737

(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

2711

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.

2716

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

2719

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

2719

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

2719

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

777

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

798

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

1013

In what scenarios would this case be triggered?

2593

2596

2718

Fit this comment to 80 characters long lines.

2723

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

2724

spell out the auto

2766

2767

2783

2845

spell out the auto

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

62 ↗

(On Diff #424252)

nit

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

1173 ↗

(On Diff #424252)

spell out the auto

1174 ↗

(On Diff #424252)

prefer unsigned

[mehdi_amini]

LGTM overall!

mlir/include/mlir/IR/OperationSupport.h

777

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

mlir/lib/IR/AsmPrinter.cpp

1013

Invalid IR?

2723

That work only with SmallVector AFAIK

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

1174 ↗

(On Diff #424252)

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

1151

Can you reflow this comment?

2709

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

2711

Please spell out auto here.

2715

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

2718

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

2718

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

2720

Please use punctuation for documentation.

2724

Please wrap the for in braces as well.

2730

Please add punctuation to the end of comments.

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

1174 ↗

(On Diff #424252)

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

[cpillmayer]

mlir/lib/IR/AsmPrinter.cpp

1013

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 ↗

(On Diff #424252)

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:
   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 a result.
+  bool shouldPrintOperationUsers() 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 results.
+  bool printOperationUsers : 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> printOperationUsers{
+      "mlir-print-operation-users", llvm::cl::init(false),
+      llvm::cl::desc("Print users of an operation 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), printOperationUsers(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;
+  printOperationUsers = clOptions->printOperationUsers;
 }
 
 /// 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) {

 /// 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 a result.
+bool OpPrintingFlags::shouldPrintOperationUsers() const {
+  return printOperationUsers;
+}
+
 /// 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 results, 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 users of operations should be printed,
+    // make sure that operations with no result have an id
+    if (printerFlags.shouldPrintOperationUsers()) {
+      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 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.shouldPrintOperationUsers())
+    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);
 }
 
+/// Print users of this operation or id of this operation if it has no result.
+void OperationPrinter::printUsersComment(Operation *op) {
+  if (!op->getNumResults() && op->getNumOperands()) {
+    os << " // id: ";
+    printOperationID(op);
+  } else if (op->getNumResults() && op->getUsers().empty()) {
+    // Indicate when there is an unused result.
+    os << " // unused";
+  } else if (!op->getUsers().empty()) {
+    os << " // " << (op->hasOneUse() ? "user" : "users") << ": ";
+    SmallPtrSet<Operation *, 8> userSet;
+    for (auto indexedUser : llvm::enumerate(op->getUsers())) {
+      Operation *user = indexedUser.value();
+      if (!userSet.insert(user).second)
+        continue;
+      printUserIDs(user, indexedUser.index());
+    }
+  }
+}
+
+/// Print users of this block arg.
+void OperationPrinter::printUsersComment(BlockArgument &arg) {
+  SmallPtrSet<Operation *, 8> userSet;
+  os << " used by ";
+  for (auto indexedUser : llvm::enumerate(arg.getUsers())) {
+    Operation *user = indexedUser.value();
+    if (!userSet.insert(user).second)
+      continue;
+    printUserIDs(user, indexedUser.index());
+  }
+}
+
+/// Prints either the ids of the result values or the id of the operation if the
+/// operation has no results.
+void OperationPrinter::printUserIDs(Operation *user, bool prefixComma) {
+  if (prefixComma)
+    os << ", ";
+
+  if (!user->getNumResults())
+    printOperationID(user);
+  else
+    // just print the id once from the first result (instead of: %1#0, %1#1,
+    // ...)
+    printValueID(user->getResult(0), false);
+}
+
 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 << ')';
 
   // For terminators, print the list of successors and their operands.
   if (op->getNumSuccessors() != 0) {
     os << '[';
     interleaveComma(op->getSuccessors(),
                     [&](Block *successor) { printBlockName(successor); });
     os << ']';
   }
 
   // Print regions.
   if (op->getNumRegions() != 0) {
     os << " (";
     interleaveComma(op->getRegions(), [&](Region &region) {
       printRegion(region, /*printEntryBlockArgs=*/true,
                   /*printBlockTerminators=*/true, /*printEmptyBlock=*/true);
     });
     os << ')';
   }
 
   auto attrs = op->getAttrs();
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
   os << " : ";
   printFunctionalType(op);
 }
 
 void OperationPrinter::printBlockName(Block *block) {
if (!block->getParent()) {
       os << "  // " << predIDs.size() << " preds: ";
 
       interleaveComma(predIDs, [&](BlockInfo pred) { os << pred.name; });
     }
     os << newLine;
   }
 
   currentIndent += indentWidth;
+
+  if (printerFlags.shouldPrintOperationUsers()) {
+    for (auto arg : block->getArguments()) {
+      if (arg.getUsers().empty())
+        continue;
+      os.indent(currentIndent);
+      os << "// ";
+      printValueID(arg);
+      printUsersComment(arg);
+      os << newLine;
+    }
+  }
+
   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())