[mlir][core] Add option to print users of an operation as a comment

[cpillmayer]

Summary

This capability in the AssemblyWriter allows to print users of an operation as a comment as proposed in git issue #53286.
To be able to refer to void type users such as store or ret, an option to generate slot numbers for void instructions is added.

[rriddle]

There may be some confusion happening here. It looks like you intended to add support for MLIR (given the commit title and bug reference), but the files in this commit are adding support to LLVM instead. MLIR and LLVM have different underlying infrastructure. (though this may still be a welcome addition for LLVM users)

[cpillmayer]

In D123959#3457517, @rriddle wrote:

There may be some confusion happening here. It looks like you intended to add support for MLIR (given the commit title and bug reference), but the files in this commit are adding support to LLVM instead. MLIR and LLVM have different underlying infrastructure. (though this may still be a welcome addition for LLVM users)

@rriddle You are right, sorry about that. I will close this patch and reopen it with the correct description.

llvm/include/llvm/IR/BasicBlock.h

public:
   BasicBlock *getUniqueSuccessor() {
     return const_cast<BasicBlock *>(
                    static_cast<const BasicBlock *>(this)->getUniqueSuccessor());
   }
 
   /// Print the basic block to an output stream with an optional
   /// AssemblyAnnotationWriter.
   void print(raw_ostream &OS, AssemblyAnnotationWriter *AAW = nullptr,
-             bool ShouldPreserveUseListOrder = false,
-             bool IsForDebug = false) const;
+             bool ShouldPreserveUseListOrder = false, bool IsForDebug = false,
+             bool AnnotateUsers = false) const;
 
   //===--------------------------------------------------------------------===//
   /// Instruction iterator methods
   ///
   inline iterator                begin()       { return InstList.begin(); }
   inline const_iterator          begin() const { return InstList.begin(); }
   inline iterator                end  ()       { return InstList.end();   }
   inline const_iterator          end  () const { return InstList.end();   }

llvm/include/llvm/IR/Function.h

/// @}
 
   /// Get the prologue data associated with this function.
   Constant *getPrologueData() const;
   void setPrologueData(Constant *PrologueData);
 
   /// Print the function to an output stream with an optional
   /// AssemblyAnnotationWriter.
   void print(raw_ostream &OS, AssemblyAnnotationWriter *AAW = nullptr,
-             bool ShouldPreserveUseListOrder = false,
-             bool IsForDebug = false) const;
+             bool ShouldPreserveUseListOrder = false, bool IsForDebug = false,
+             bool AnnotateUsers = false) const;
 
   /// viewCFG - This function is meant for use from the debugger.  You can just
   /// say 'call F->viewCFG()' and a ghostview window should pop up from the
   /// program, displaying the CFG of the current function with the code for each
   /// basic block inside.  This depends on there being a 'dot' and 'gv' program
   /// in your path.
   ///
   void viewCFG() const;

llvm/include/llvm/IR/Module.h

 /// @name Utility functions for printing and dumping Module objects
 /// @{
 
   /// Print the module to an output stream with an optional
   /// AssemblyAnnotationWriter.  If \c ShouldPreserveUseListOrder, then include
   /// uselistorder directives so that use-lists can be recreated when reading
   /// the assembly.
   void print(raw_ostream &OS, AssemblyAnnotationWriter *AAW,
-             bool ShouldPreserveUseListOrder = false,
-             bool IsForDebug = false) const;
+             bool ShouldPreserveUseListOrder = false, bool IsForDebug = false,
+             bool AnnotateUsers = false) const;
 
   /// Dump the module to stderr (for debugging).
   void dump() const;
 
   /// This function causes all the subinstructions to "let go" of all references
   /// that they are maintaining.  This allows one to 'delete' a whole class at
   /// a time, even though there may be circular references... first all
   /// references are dropped, and all use counts go to zero.  Then everything

llvm/lib/IR/AsmWriter.cpp

private:
   /// TheModule - The module for which we are holding slot numbers.
   const Module* TheModule;
 
   /// TheFunction - The function for which we are holding slot numbers.
   const Function* TheFunction = nullptr;
   bool FunctionProcessed = false;
   bool ShouldInitializeAllMetadata;
 
+  /// Assign slot numbers for void type instructions
+  bool IncludeVoidInstructions;
+
   std::function<void(AbstractSlotTrackerStorage *, const Module *, bool)>
       ProcessModuleHookFn;
   std::function<void(AbstractSlotTrackerStorage *, const Function *, bool)>
       ProcessFunctionHookFn;
 
   /// The summary index for which we are holding slot numbers.
   const ModuleSummaryIndex *TheIndex = nullptr;
 

 
 public:
   /// Construct from a module.
   ///
   /// If \c ShouldInitializeAllMetadata, initializes all metadata in all
   /// functions, giving correct numbering for metadata referenced only from
   /// within a function (even if no functions have been initialized).
   explicit SlotTracker(const Module *M,
-                       bool ShouldInitializeAllMetadata = false);
+                       bool ShouldInitializeAllMetadata = false,
+                       bool IncludeVoidInstructions = false);
 
   /// Construct from a function, starting out in incorp state.
   ///
   /// If \c ShouldInitializeAllMetadata, initializes all metadata in all
   /// functions, giving correct numbering for metadata referenced only from
   /// within a function (even if no functions have been initialized).
   explicit SlotTracker(const Function *F,
-                       bool ShouldInitializeAllMetadata = false);
+                       bool ShouldInitializeAllMetadata = false,
+                       bool IncludeVoidInstructions = false);
 
   /// Construct from a module summary index.
   explicit SlotTracker(const ModuleSummaryIndex *Index);
 
   SlotTracker(const SlotTracker &) = delete;
   SlotTracker &operator=(const SlotTracker &) = delete;
 
   ~SlotTracker() = default;

 #if 0
 #define ST_DEBUG(X) dbgs() << X
 #else
 #define ST_DEBUG(X)
 #endif
 
 // Module level constructor. Causes the contents of the Module (sans functions)
 // to be added to the slot table.
-SlotTracker::SlotTracker(const Module *M, bool ShouldInitializeAllMetadata)
-    : TheModule(M), ShouldInitializeAllMetadata(ShouldInitializeAllMetadata) {}
+SlotTracker::SlotTracker(const Module *M, bool ShouldInitializeAllMetadata,
+                         bool IncludeVoidInstructions)
+    : TheModule(M), ShouldInitializeAllMetadata(ShouldInitializeAllMetadata),
+      IncludeVoidInstructions(IncludeVoidInstructions) {}
 
 // Function level constructor. Causes the contents of the Module and the one
 // function provided to be added to the slot table.
-SlotTracker::SlotTracker(const Function *F, bool ShouldInitializeAllMetadata)
+SlotTracker::SlotTracker(const Function *F, bool ShouldInitializeAllMetadata,
+                         bool IncludeVoidInstructions)
     : TheModule(F ? F->getParent() : nullptr), TheFunction(F),
-      ShouldInitializeAllMetadata(ShouldInitializeAllMetadata) {}
+      ShouldInitializeAllMetadata(ShouldInitializeAllMetadata),
+      IncludeVoidInstructions(IncludeVoidInstructions) {}
 
 SlotTracker::SlotTracker(const ModuleSummaryIndex *Index)
-    : TheModule(nullptr), ShouldInitializeAllMetadata(false), TheIndex(Index) {}
+    : TheModule(nullptr), ShouldInitializeAllMetadata(false),
+      IncludeVoidInstructions(false), TheIndex(Index) {}
 
 inline void SlotTracker::initializeIfNeeded() {
   if (TheModule) {
     processModule();
     TheModule = nullptr; ///< Prevent re-processing next time we're called.
   }
 
   if (TheFunction && !FunctionProcessed)
void SlotTracker::processFunction() {
   ST_DEBUG("Inserting Instructions:\n");
 
   // Add all of the basic blocks and instructions with no names.
   for (auto &BB : *TheFunction) {
     if (!BB.hasName())
       CreateFunctionSlot(&BB);
 
     for (auto &I : BB) {
-      if (!I.getType()->isVoidTy() && !I.hasName())
+      if ((!I.getType()->isVoidTy() || IncludeVoidInstructions) && !I.hasName())
         CreateFunctionSlot(&I);
 
       // We allow direct calls to any llvm.foo function here, because the
       // target may not be linked into the optimizer.
       if (const auto *Call = dyn_cast<CallBase>(&I)) {
         // Add all the call attributes to the table.
         AttributeSet Attrs = Call->getAttributes().getFnAttrs();
         if (Attrs.hasAttributes())
class AssemblyWriter {
   const ModuleSummaryIndex *TheIndex = nullptr;
   std::unique_ptr<SlotTracker> SlotTrackerStorage;
   SlotTracker &Machine;
   TypePrinting TypePrinter;
   AssemblyAnnotationWriter *AnnotationWriter = nullptr;
   SetVector<const Comdat *> Comdats;
   bool IsForDebug;
   bool ShouldPreserveUseListOrder;
+  bool AnnotateUsers;
   UseListOrderMap UseListOrders;
   SmallVector<StringRef, 8> MDNames;
   /// Synchronization scope names registered with LLVMContext.
   SmallVector<StringRef, 8> SSNs;
   DenseMap<const GlobalValueSummary *, GlobalValue::GUID> SummaryToGUIDMap;
 
 public:
   /// Construct an AssemblyWriter with an external SlotTracker
   AssemblyWriter(formatted_raw_ostream &o, SlotTracker &Mac, const Module *M,
                  AssemblyAnnotationWriter *AAW, bool IsForDebug,
-                 bool ShouldPreserveUseListOrder = false);
+                 bool ShouldPreserveUseListOrder = false,
+                 bool AnnotateUsers = false);
 
   AssemblyWriter(formatted_raw_ostream &o, SlotTracker &Mac,
                  const ModuleSummaryIndex *Index, bool IsForDebug);
 
   AsmWriterContext getContext() {
     return AsmWriterContext(&TypePrinter, &Machine, TheModule);
   }
 
private:
   void printMetadataAttachments(
       const SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs,
       StringRef Separator);
 
   // printInfoComment - Print a little comment after the instruction indicating
   // which slot it occupies.
   void printInfoComment(const Value &V);
 
+  // printUsersComment - Print a comment indicating the users of the result of
+  // this value
+  void printUsersComment(const Value &V);
+
   // printGCRelocateComment - print comment after call to the gc.relocate
   // intrinsic indicating base and derived pointer names.
   void printGCRelocateComment(const GCRelocateInst &Relocate);
 };
 
 } // end anonymous namespace
 
 AssemblyWriter::AssemblyWriter(formatted_raw_ostream &o, SlotTracker &Mac,
                                const Module *M, AssemblyAnnotationWriter *AAW,
-                               bool IsForDebug, bool ShouldPreserveUseListOrder)
+                               bool IsForDebug, bool ShouldPreserveUseListOrder,
+                               bool AnnotateUsers)
     : Out(o), TheModule(M), Machine(Mac), TypePrinter(M), AnnotationWriter(AAW),
       IsForDebug(IsForDebug),
-      ShouldPreserveUseListOrder(ShouldPreserveUseListOrder) {
+      ShouldPreserveUseListOrder(ShouldPreserveUseListOrder),
+      AnnotateUsers(AnnotateUsers) {
   if (!TheModule)
     return;
   for (const GlobalObject &GO : TheModule->global_objects())
     if (const Comdat *C = GO.getComdat())
       Comdats.insert(C);
 }
 
 AssemblyWriter::AssemblyWriter(formatted_raw_ostream &o, SlotTracker &Mac,
                                const ModuleSummaryIndex *Index, bool IsForDebug)
     : Out(o), TheIndex(Index), Machine(Mac), TypePrinter(/*Module=*/nullptr),
-      IsForDebug(IsForDebug), ShouldPreserveUseListOrder(false) {}
+      IsForDebug(IsForDebug), ShouldPreserveUseListOrder(false),
+      AnnotateUsers(false) {}
 
 void AssemblyWriter::writeOperand(const Value *Operand, bool PrintType) {
   if (!Operand) {
     Out << "<null operand!>";
     return;
   }
   if (PrintType) {
     TypePrinter.print(Operand->getType(), Out);
void AssemblyWriter::printFunction(const Function *F) {
   if (F->isDeclaration()) {
     Out << '\n';
   } else {
     SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
     F->getAllMetadata(MDs);
     printMetadataAttachments(MDs, " ");
 
     Out << " {";
+
+    // Print users comments for arguments
+    if (AnnotateUsers) {
+      for (auto &Arg : F->args()) {
+        Out << "\n";
+        int SlotNum = Machine.getLocalSlot(&Arg);
+        Out << "; %" << SlotNum;
+        printUsersComment(Arg);
+      }
+    }
+
     // Output all of the function's basic blocks.
     for (const BasicBlock &BB : *F)
       printBasicBlock(&BB);
 
     // Output the function's use-lists.
     printUseLists(F);
 
     Out << "}\n";
void AssemblyWriter::printInstruction(const Instruction &I) {
 
   // Print Metadata info.
   SmallVector<std::pair<unsigned, MDNode *>, 4> InstMD;
   I.getAllMetadata(InstMD);
   printMetadataAttachments(InstMD, ", ");
 
   // Print a nice comment.
   printInfoComment(I);
+
+  // Print out slot number for void instructions and users otherwise
+  if (AnnotateUsers) {
+    if (I.getType()->isVoidTy()) {
+      Out.PadToColumn(50);
+      Out << "; id = %";
+      if (I.hasName())
+        Out << I.getName();
+      else
+        Out << Machine.getLocalSlot(&I);
+    } else {
+      printUsersComment(I);
+    }
+  }
+}
+
+void AssemblyWriter::printUsersComment(const Value &V) {
+
+  if (V.user_empty())
+    return;
+
+  Out.PadToColumn(50);
+  Out << "; " << (V.hasOneUser() ? "user" : "users") << " = ";
+  for (auto U = V.user_begin(); U != V.user_end(); U++) {
+    if (U != V.user_begin())
+      Out << ", ";
+
+    Out << "%";
+    if (U->hasName()) {
+      Out << U->getName();
+    } else {
+      Out << Machine.getLocalSlot(*U);
+    }
+  }
 }
 
 void AssemblyWriter::printMetadataAttachments(
     const SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs,
     StringRef Separator) {
   if (MDs.empty())
     return;
 
for (const auto &Pair : It->second)
     printUseListOrder(Pair.first, Pair.second);
 }
 
 //===----------------------------------------------------------------------===//
 //                       External Interface declarations
 //===----------------------------------------------------------------------===//
 
 void Function::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
-                     bool ShouldPreserveUseListOrder,
-                     bool IsForDebug) const {
-  SlotTracker SlotTable(this->getParent());
+                     bool ShouldPreserveUseListOrder, bool IsForDebug,
+                     bool AnnotateUsers) const {
+  SlotTracker SlotTable(this->getParent(), false, AnnotateUsers);
   formatted_raw_ostream OS(ROS);
-  AssemblyWriter W(OS, SlotTable, this->getParent(), AAW,
-                   IsForDebug,
-                   ShouldPreserveUseListOrder);
+  AssemblyWriter W(OS, SlotTable, this->getParent(), AAW, IsForDebug,
+                   ShouldPreserveUseListOrder, AnnotateUsers);
   W.printFunction(this);
 }
 
 void BasicBlock::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
-                     bool ShouldPreserveUseListOrder,
-                     bool IsForDebug) const {
-  SlotTracker SlotTable(this->getParent());
+                       bool ShouldPreserveUseListOrder, bool IsForDebug,
+                       bool AnnotateUsers) const {
+  SlotTracker SlotTable(this->getParent(), false, AnnotateUsers);
   formatted_raw_ostream OS(ROS);
-  AssemblyWriter W(OS, SlotTable, this->getModule(), AAW,
-                   IsForDebug,
-                   ShouldPreserveUseListOrder);
+  AssemblyWriter W(OS, SlotTable, this->getModule(), AAW, IsForDebug,
+                   ShouldPreserveUseListOrder, AnnotateUsers);
   W.printBasicBlock(this);
 }
 
 void Module::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
-                   bool ShouldPreserveUseListOrder, bool IsForDebug) const {
-  SlotTracker SlotTable(this);
+                   bool ShouldPreserveUseListOrder, bool IsForDebug,
+                   bool AnnotateUsers) const {
+  SlotTracker SlotTable(this, false, AnnotateUsers);
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this, AAW, IsForDebug,
-                   ShouldPreserveUseListOrder);
+                   ShouldPreserveUseListOrder, AnnotateUsers);
   W.printModule(this);
 }
 
 void NamedMDNode::print(raw_ostream &ROS, bool IsForDebug) const {
   SlotTracker SlotTable(getParent());
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, getParent(), nullptr, IsForDebug);
   W.printNamedMDNode(this);