[DebugInfo] Enable the capability of attaching metadata to BasicBlock.

[HsiangKai]

Summary

In order to attach debug info to labels. We need to enable the capability of attaching metadata to BasicBlock. If we have following LLVM IR fragment,

top:

%0 = load i32, i32* %a.addr, align 4

we could attach debug metadata to label 'top' as follow

top: !10

%0 = load i32, i32* %a.addr, align 4

!10 = DILabel(name: top, ...)

[HsiangKai]

Add DILabel to SlotTracker in IR AsmWriter.

[chenwj]

I don't understand the part of BasicBlock, the others LGTM.

include/llvm/IR/BasicBlock.h

66

Would you mind elaborating what the enum, RefCountBits and RefCountMask are for?

lib/IR/AsmWriter.cpp

2930

If I understand correctly, BB->getAllMetadata() should return MD of label? InstMD is kind of confusing here. It makes me think you want to get MD of instructions within the basic block. Maybe LabelMD would be better?

[aprantl]

I replied to the thread on llvm-dev. We also need a plan for how to update this metadata in the inliner.

[HsiangKai]

As discussed in llvm-dev mailing list, we should attach label metadata in intrinsic instead of basic block.

include/llvm/IR/BasicBlock.h

 class CallInst;
 class Function;
 class LandingPadInst;
 class LLVMContext;
 class Module;
 class PHINode;
 class TerminatorInst;
 class ValueSymbolTable;
+class DILabel;
 
 /// \brief LLVM Basic Block Representation
 ///
 /// This represents a single basic block in LLVM. A basic block is simply a
 /// container of instructions that execute sequentially. Basic blocks are Values
 /// because they are referenced by instructions such as branches and switch
 /// tables. The type of a BasicBlock is "Type::LabelTy" because the basic block
 /// represents a label to which a branch can jump.
 ///
 /// A well formed basic block is formed of a list of non-terminating
 /// instructions followed by a single TerminatorInst instruction.
 /// TerminatorInst's may not occur in the middle of basic blocks, and must
 /// terminate the blocks. The BasicBlock class allows malformed basic blocks to
 /// occur because it may be useful in the intermediate stage of constructing or
 /// modifying a program. However, the verifier will ensure that basic blocks
 /// are "well formed".
 class BasicBlock final : public Value, // Basic blocks are data objects also
                          public ilist_node_with_parent<BasicBlock, Function> {
 public:
   using InstListType = SymbolTableList<Instruction>;
 
 private:
+  enum {
+    LastRefCountBit = 3,
+    /// This is a bit stored in the SubClassData field which indicates whether
+    /// this basic block has metadata attached to it or not.
+    HasMetadataBit,
+  };
+  static const unsigned RefCountBits = LastRefCountBit + 1;
+  static const unsigned RefCountMask = (1 << RefCountBits) - 1;
+
   friend class BlockAddress;
   friend class SymbolTableListTraits<BasicBlock>;
 
   InstListType InstList;
   Function *Parent;
 
   void setParent(Function *parent);
 
public:
   /// keeping loop information consistent, use the SplitBlock utility function.
   BasicBlock *splitBasicBlock(iterator I, const Twine &BBName = "");
   BasicBlock *splitBasicBlock(Instruction *I, const Twine &BBName = "") {
     return splitBasicBlock(I->getIterator(), BBName);
   }
 
   /// \brief Returns true if there are any uses of this basic block other than
   /// direct branches, switches, etc. to it.
-  bool hasAddressTaken() const { return getSubclassDataFromValue() != 0; }
+  bool hasAddressTaken() const {
+    return (getSubclassDataFromValue() & RefCountMask) != 0;
+  }
 
   /// \brief Update all phi nodes in this basic block's successors to refer to
   /// basic block \p New instead of to it.
   void replaceSuccessorsPhiUsesWith(BasicBlock *New);
 
   /// \brief Return true if this basic block is an exception handling block.
   bool isEHPad() const { return getFirstNonPHI()->isEHPad(); }
 
return const_cast<LandingPadInst *>(
                     static_cast<const BasicBlock *>(this)->getLandingPadInst());
   }
 
   /// \brief Return true if it is legal to hoist instructions into this block.
   bool isLegalToHoistInto() const;
 
   Optional<uint64_t> getIrrLoopHeaderWeight() const;
 
+  //===--------------------------------------------------------------------===//
+  // Metadata manipulation.
+  //===--------------------------------------------------------------------===//
+
+  /// Return true if this basic block has any metadata attached to it.
+  bool hasMetadata() const { return hasMetadataHashEntry(); }
+
+  /// Get the metadata of given kind attached to this Basic Block.
+  /// If the metadata is not found then return null.
+  MDNode *getMetadata(unsigned KindID) const;
+  MDNode *getMetadata(StringRef Kind) const;
+
+  /// Get all metadata attached to this BasicBlock. The first element of each
+  /// pair returned is the KindID, the second element is the metadata value.
+  /// This list is returned sorted by the KindID.
+  void
+  getAllMetadata(SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) const;
+
+  /// Set the metadata of the specified kind to the specified node. This updates
+  /// or replaces metadata if already present, or removes it if Node is null.
+  void setMetadata(unsigned KindID, MDNode *Node);
+  void setMetadata(StringRef Kind, MDNode *Node);
+
+  /// \brief Set the attached label.
+  ///
+  /// Calls \a setMetadata() with \a LLVMContext::MD_dbg.
+  void setLabel(DILabel *L);
+
+  /// \brief Get the attached label.
+  ///
+  /// Calls \a getMetadata() with \a LLVMContext::MD_dbg and casts the result
+  /// to \a DILabel.
+  DILabel *getLabel() const;
+
 private:
+  /// Return true if we have an entry in the on-the-side metadata hash.
+  bool hasMetadataHashEntry() const {
+    return (getSubclassDataFromValue() & HasMetadataBit) != 0;
+  }
+
+  void setHasMetadataHashEntry(bool V) {
+    setValueSubclassData((getSubclassDataFromValue() & ~HasMetadataBit) |
+                         (V ? HasMetadataBit : 0));
+  }
+
   /// \brief Increment the internal refcount of the number of BlockAddresses
   /// referencing this BasicBlock by \p Amt.
   ///
   /// This is almost always 0, sometimes one possibly, but almost never 2, and
   /// inconceivably 3 or more.
   void AdjustBlockAddressRefCount(int Amt) {
-    setValueSubclassData(getSubclassDataFromValue()+Amt);
-    assert((int)(signed char)getSubclassDataFromValue() >= 0 &&
+    auto RefCount = getSubclassDataFromValue() & RefCountMask;
+    setValueSubclassData((getSubclassDataFromValue() & ~RefCountMask)
+                   | (RefCount + Amt));
+    assert((int)(signed char)(getSubclassDataFromValue() & RefCountMask) >= 0 &&
            "Refcount wrap-around");
   }
 
   /// \brief Shadow Value::setValueSubclassData with a private forwarding method
   /// so that any future subclasses cannot accidentally use it.
   void setValueSubclassData(unsigned short D) {
     Value::setValueSubclassData(D);
   }
 };
 
 // Create wrappers for C Binding types (see CBindingWrapping.h).
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(BasicBlock, LLVMBasicBlockRef)
 
 } // end namespace llvm
 
 #endif // LLVM_IR_BASICBLOCK_H

lib/IR/AsmWriter.cpp

private:
   void processFunction();
 
   /// Add the metadata directly attached to a GlobalObject.
   void processGlobalObjectMetadata(const GlobalObject &GO);
 
   /// Add all of the metadata from a function.
   void processFunctionMetadata(const Function &F);
 
+  /// Add all of the metadata from a basic block.
+  void processBasicBlockMetadata(const BasicBlock &BB);
+
   /// Add all of the metadata from an instruction.
   void processInstructionMetadata(const Instruction &I);
 };
 
 } // end namespace llvm
 
 ModuleSlotTracker::ModuleSlotTracker(SlotTracker &Machine, const Module *M,
                                      const Function *F)
void SlotTracker::processGlobalObjectMetadata(const GlobalObject &GO) {
   GO.getAllMetadata(MDs);
   for (auto &MD : MDs)
     CreateMetadataSlot(MD.second);
 }
 
 void SlotTracker::processFunctionMetadata(const Function &F) {
   processGlobalObjectMetadata(F);
   for (auto &BB : F) {
+    processBasicBlockMetadata(BB);
     for (auto &I : BB)
       processInstructionMetadata(I);
   }
 }
 
+void SlotTracker::processBasicBlockMetadata(const BasicBlock &BB) {
+  if (MDNode *N = BB.getLabel())
+    CreateMetadataSlot(N);
+}
+
 void SlotTracker::processInstructionMetadata(const Instruction &I) {
   // Process metadata used directly by intrinsics.
   if (const CallInst *CI = dyn_cast<CallInst>(&I))
     if (Function *F = CI->getCalledFunction())
       if (F->isIntrinsic())
         for (auto &Op : I.operands())
           if (auto *V = dyn_cast_or_null<MetadataAsValue>(Op))
             if (MDNode *N = dyn_cast<MDNode>(V->getMetadata()))
if (BB->hasName()) {              // Print out the label if it exists...
     Out << "\n; <label>:";
     int Slot = Machine.getLocalSlot(BB);
     if (Slot != -1)
       Out << Slot << ":";
     else
       Out << "<badref>";
   }
 
+  // Print Metadata info.
+  SmallVector<std::pair<unsigned, MDNode *>, 4> InstMD;
+  BB->getAllMetadata(InstMD);
+  printMetadataAttachments(InstMD, " ");
+
   if (!BB->getParent()) {
     Out.PadToColumn(50);
     Out << "; Error: Block without parent!";
   } else if (BB != &BB->getParent()->getEntryBlock()) {  // Not the entry block?
     // Output predecessors for the block.
     Out.PadToColumn(50);
     Out << ";";
     const_pred_iterator PI = pred_begin(BB), PE = pred_end(BB);

lib/IR/LLVMContextImpl.h

#include "llvm/IR/Metadata.def"
   ValueHandlesTy ValueHandles;
   
   /// CustomMDKindNames - Map to hold the metadata string to ID mapping.
   StringMap<unsigned> CustomMDKindNames;
 
   /// Collection of per-instruction metadata used in this context.
   DenseMap<const Instruction *, MDAttachmentMap> InstructionMetadata;
 
+  /// Collection of per-basicblock metadata used in this context.
+  DenseMap<const BasicBlock *, MDAttachmentMap> BasicBlockMetadata;
+
   /// Collection of per-GlobalObject metadata used in this context.
   DenseMap<const GlobalObject *, MDGlobalAttachmentMap> GlobalObjectMetadata;
 
   /// Collection of per-GlobalObject sections used in this context.
   DenseMap<const GlobalObject *, StringRef> GlobalObjectSections;
 
   /// Stable collection of section strings.
   StringSet<> SectionStrings;

lib/IR/Metadata.cpp

 
 void GlobalVariable::getDebugInfo(
     SmallVectorImpl<DIGlobalVariableExpression *> &GVs) const {
   SmallVector<MDNode *, 1> MDs;
   getMetadata(LLVMContext::MD_dbg, MDs);
   for (MDNode *MD : MDs)
     GVs.push_back(cast<DIGlobalVariableExpression>(MD));
 }
+
+MDNode *BasicBlock::getMetadata(unsigned KindID) const {
+  if (!hasMetadata())
+    return nullptr;
+  auto &Info = getContext().pImpl->BasicBlockMetadata[this];
+  assert(!Info.empty() && "bit out of sync with hash table");
+
+  return Info.lookup(KindID);
+}
+
+MDNode *BasicBlock::getMetadata(StringRef Kind) const {
+  return getMetadata(getContext().getMDKindID(Kind));
+}
+
+void BasicBlock::setMetadata(unsigned KindID, MDNode *Node) {
+  if (!Node && !hasMetadata())
+    return;
+
+  // Handle the case when we're adding/updating metadata on an basic block.
+  if (Node) {
+    auto &Info = getContext().pImpl->BasicBlockMetadata[this];
+    assert(!Info.empty() == hasMetadataHashEntry() &&
+           "HasMetadata bit is wonked");
+    if (Info.empty())
+      setHasMetadataHashEntry(true);
+    Info.set(KindID, *Node);
+    return;
+  }
+
+  // Otherwise, we're removing metadata from an basic block.
+  assert((hasMetadataHashEntry() ==
+          (getContext().pImpl->BasicBlockMetadata.count(this) > 0)) &&
+         "HasMetadata bit out of date!");
+  if (!hasMetadataHashEntry())
+    return; // Nothing to remove!
+  auto &Info = getContext().pImpl->BasicBlockMetadata[this];
+
+  // Handle removal of an existing value.
+  Info.erase(KindID);
+
+  if (!Info.empty())
+    return;
+
+  getContext().pImpl->BasicBlockMetadata.erase(this);
+  setHasMetadataHashEntry(false);
+}
+
+void BasicBlock::setMetadata(StringRef Kind, MDNode *Node) {
+  if (!Node && !hasMetadata())
+    return;
+  setMetadata(getContext().getMDKindID(Kind), Node);
+}
+
+void BasicBlock::getAllMetadata(
+    SmallVectorImpl<std::pair<unsigned, MDNode *>> &Result) const {
+  Result.clear();
+
+  if (!hasMetadata())
+    return;
+
+  assert(hasMetadataHashEntry() &&
+         getContext().pImpl->BasicBlockMetadata.count(this) &&
+         "Shouldn't have called this");
+  const auto &Info = getContext().pImpl->BasicBlockMetadata.find(this)->second;
+  assert(!Info.empty() && "Shouldn't have called this");
+  Info.getAll(Result);
+}
+
+void BasicBlock::setLabel(DILabel *L) {
+  setMetadata(LLVMContext::MD_dbg, L);
+}
+
+DILabel *BasicBlock::getLabel() const {
+  return cast_or_null<DILabel>(getMetadata(LLVMContext::MD_dbg));
+}