[llvm][AsmPrinter] Add direct calls to callgraph section (#155706)

Extend CallGraphSection to include metadata about direct calls. This
simplifies the design of tools that must parse .callgraph section to not
require dependency on MC layer.

Author: Prabhuk

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/StaticDataProfileInfo.h"
@@ -192,28 +193,28 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
 
   /// Store symbols and type identifiers used to create callgraph section
   /// entries related to a function.
-  struct FunctionInfo {
+  struct FunctionCallGraphInfo {
     /// Numeric type identifier used in callgraph section for indirect calls
     /// and targets.
     using CGTypeId = uint64_t;
 
-    /// Enumeration of function kinds, and their mapping to function kind values
-    /// stored in callgraph section entries.
-    /// Must match the enum in llvm/tools/llvm-objdump/llvm-objdump.cpp.
-    enum class FunctionKind : uint64_t {
-      /// Function cannot be target to indirect calls.
-      NOT_INDIRECT_TARGET = 0,
-
-      /// Function may be target to indirect calls but its type id is unknown.
-      INDIRECT_TARGET_UNKNOWN_TID = 1,
-
-      /// Function may be target to indirect calls and its type id is known.
-      INDIRECT_TARGET_KNOWN_TID = 2,
-    };
-
     /// Map type identifiers to callsite labels. Labels are generated for each
     /// indirect callsite in the function.
     SmallVector<std::pair<CGTypeId, MCSymbol *>> CallSiteLabels;
+    SmallSet<MCSymbol *, 4> DirectCallees;
+  };
+
+  /// Enumeration of function kinds, and their mapping to function kind values
+  /// stored in callgraph section entries.
+  enum class FunctionKind : uint64_t {
+    /// Function cannot be target to indirect calls.
+    NOT_INDIRECT_TARGET = 0,
+
+    /// Function may be target to indirect calls but its type id is unknown.
+    INDIRECT_TARGET_UNKNOWN_TID = 1,
+
+    /// Function may be target to indirect calls and its type id is known.
+    INDIRECT_TARGET_KNOWN_TID = 2,
   };
 
   enum CallGraphSectionFormatVersion : uint64_t {
@@ -385,10 +386,11 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
   /// are available. Returns empty string otherwise.
   StringRef getConstantSectionSuffix(const Constant *C) const;
 
-  /// Generate and emit labels for callees of the indirect callsites which will
-  /// be used to populate the .callgraph section.
-  void emitIndirectCalleeLabels(
-      FunctionInfo &FuncInfo,
+  /// Iff MI is an indirect call, generate and emit a label after the callsites
+  /// which will be used to populate the .callgraph section. For direct
+  /// callsites add the callee symbol to direct callsites list of FuncCGInfo.
+  void handleCallsiteForCallgraph(
+      FunctionCallGraphInfo &FuncCGInfo,
       const MachineFunction::CallSiteInfoMap &CallSitesInfoMap,
       const MachineInstr &MI);
 
@@ -479,7 +481,8 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
   void emitKCFITrapEntry(const MachineFunction &MF, const MCSymbol *Symbol);
   virtual void emitKCFITypeId(const MachineFunction &MF);
 
-  void emitCallGraphSection(const MachineFunction &MF, FunctionInfo &FuncInfo);
+  void emitCallGraphSection(const MachineFunction &MF,
+                            FunctionCallGraphInfo &FuncCGInfo);
 
   void emitPseudoProbe(const MachineInstr &MI);
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -78,6 +78,7 @@
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMRemarkStreamer.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Metadata.h"
@@ -1673,7 +1674,7 @@ static ConstantInt *extractNumericCGTypeId(const Function &F) {
 
 /// Emits .callgraph section.
 void AsmPrinter::emitCallGraphSection(const MachineFunction &MF,
-                                      FunctionInfo &FuncInfo) {
+                                      FunctionCallGraphInfo &FuncCGInfo) {
   if (!MF.getTarget().Options.EmitCallGraphSection)
     return;
 
@@ -1712,27 +1713,34 @@ void AsmPrinter::emitCallGraphSection(const MachineFunction &MF,
   // Emit function kind, and type id if available.
   if (!IsIndirectTarget) {
     OutStreamer->emitInt64(
-        static_cast<uint64_t>(FunctionInfo::FunctionKind::NOT_INDIRECT_TARGET));
+        static_cast<uint64_t>(FunctionKind::NOT_INDIRECT_TARGET));
   } else {
     if (const auto *TypeId = extractNumericCGTypeId(F)) {
-      OutStreamer->emitInt64(static_cast<uint64_t>(
-          FunctionInfo::FunctionKind::INDIRECT_TARGET_KNOWN_TID));
+      OutStreamer->emitInt64(
+          static_cast<uint64_t>(FunctionKind::INDIRECT_TARGET_KNOWN_TID));
       OutStreamer->emitInt64(TypeId->getZExtValue());
     } else {
-      OutStreamer->emitInt64(static_cast<uint64_t>(
-          FunctionInfo::FunctionKind::INDIRECT_TARGET_UNKNOWN_TID));
+      OutStreamer->emitInt64(
+          static_cast<uint64_t>(FunctionKind::INDIRECT_TARGET_UNKNOWN_TID));
     }
   }
 
   // Emit callsite labels, where each element is a pair of type id and
   // indirect callsite pc.
-  const auto &CallSiteLabels = FuncInfo.CallSiteLabels;
+  const auto &CallSiteLabels = FuncCGInfo.CallSiteLabels;
   OutStreamer->emitInt64(CallSiteLabels.size());
   for (const auto &[TypeId, Label] : CallSiteLabels) {
     OutStreamer->emitInt64(TypeId);
     OutStreamer->emitSymbolValue(Label, TM.getProgramPointerSize());
   }
-  FuncInfo.CallSiteLabels.clear();
+  FuncCGInfo.CallSiteLabels.clear();
+
+  const auto &DirectCallees = FuncCGInfo.DirectCallees;
+  OutStreamer->emitInt64(DirectCallees.size());
+  for (const auto &CalleeSymbol : DirectCallees) {
+    OutStreamer->emitSymbolValue(CalleeSymbol, TM.getProgramPointerSize());
+  }
+  FuncCGInfo.DirectCallees.clear();
 
   OutStreamer->popSection();
 }
@@ -1867,20 +1875,40 @@ static StringRef getMIMnemonic(const MachineInstr &MI, MCStreamer &Streamer) {
   return Name;
 }
 
-void AsmPrinter::emitIndirectCalleeLabels(
-    FunctionInfo &FuncInfo,
+void AsmPrinter::handleCallsiteForCallgraph(
+    FunctionCallGraphInfo &FuncCGInfo,
     const MachineFunction::CallSiteInfoMap &CallSitesInfoMap,
     const MachineInstr &MI) {
-  // Only indirect calls have type identifiers set.
+  assert(MI.isCall() &&
+         "Callsite labels are meant for call instructions only.");
+  const MachineOperand &CalleeOperand = MI.getOperand(0);
+  if (CalleeOperand.isGlobal() || CalleeOperand.isSymbol()) {
+    // Handle direct calls.
+    MCSymbol *CalleeSymbol = nullptr;
+    switch (CalleeOperand.getType()) {
+    case llvm::MachineOperand::MO_GlobalAddress:
+      CalleeSymbol = getSymbol(CalleeOperand.getGlobal());
+      break;
+    case llvm::MachineOperand::MO_ExternalSymbol:
+      CalleeSymbol = GetExternalSymbolSymbol(CalleeOperand.getSymbolName());
+      break;
+    default:
+      llvm_unreachable(
+          "Expected to only handle direct call instructions here.");
+    }
+    FuncCGInfo.DirectCallees.insert(CalleeSymbol);
+    return; // Early exit after handling the direct call instruction.
+  }
   const auto &CallSiteInfo = CallSitesInfoMap.find(&MI);
   if (CallSiteInfo == CallSitesInfoMap.end())
     return;
-
+  // Handle indirect callsite info.
+  // Only indirect calls have type identifiers set.
   for (ConstantInt *CalleeTypeId : CallSiteInfo->second.CalleeTypeIds) {
     MCSymbol *S = MF->getContext().createTempSymbol();
     OutStreamer->emitLabel(S);
     uint64_t CalleeTypeIdVal = CalleeTypeId->getZExtValue();
-    FuncInfo.CallSiteLabels.emplace_back(CalleeTypeIdVal, S);
+    FuncCGInfo.CallSiteLabels.emplace_back(CalleeTypeIdVal, S);
   }
 }
 
@@ -1930,7 +1958,7 @@ void AsmPrinter::emitFunctionBody() {
     MBBSectionRanges[MF->front().getSectionID()] =
         MBBSectionRange{CurrentFnBegin, nullptr};
 
-  FunctionInfo FuncInfo;
+  FunctionCallGraphInfo FuncCGInfo;
   const auto &CallSitesInfoMap = MF->getCallSitesInfo();
   for (auto &MBB : *MF) {
     // Print a label for the basic block.
@@ -2067,7 +2095,7 @@ void AsmPrinter::emitFunctionBody() {
         OutStreamer->emitLabel(createCallsiteEndSymbol(MBB));
 
       if (TM.Options.EmitCallGraphSection && MI.isCall())
-        emitIndirectCalleeLabels(FuncInfo, CallSitesInfoMap, MI);
+        handleCallsiteForCallgraph(FuncCGInfo, CallSitesInfoMap, MI);
 
       // If there is a post-instruction symbol, emit a label for it here.
       if (MCSymbol *S = MI.getPostInstrSymbol())
@@ -2249,7 +2277,7 @@ void AsmPrinter::emitFunctionBody() {
   emitStackSizeSection(*MF);
 
   // Emit section containing call graph metadata.
-  emitCallGraphSection(*MF, FuncInfo);
+  emitCallGraphSection(*MF, FuncCGInfo);
 
   // Emit .su file containing function stack size information.
   emitStackUsage(*MF);

llvm/test/CodeGen/X86/call-graph-section-assembly.ll

@@ -1,22 +1,34 @@
-;; Test if temporary labels are generated for each indirect callsite with a callee_type metadata.
-;; Test if the .callgraph section contains the MD5 hash of callee type ids generated from
-;; generalized type id strings.
+;; Test if temporary labels are generated for each indirect callsite.
+;; Test if the .callgraph section contains the MD5 hash of callees' type (type id)
+;; is correctly paired with its corresponding temporary label generated for indirect
+;; call sites annotated with !callee_type metadata.
+;; Test if the .callgraph section contains unique direct callees.
 
 ; RUN: llc -mtriple=x86_64-unknown-linux --call-graph-section -o - < %s | FileCheck %s
 
+declare !type !0 void @direct_foo()
+declare !type !1 i32 @direct_bar(i8)
+declare !type !2 ptr @direct_baz(ptr)
+
 ; CHECK: ball:
 ; CHECK-NEXT: [[LABEL_FUNC:\.Lfunc_begin[0-9]+]]:
 define ptr @ball() {
 entry:
+  call void @direct_foo()
   %fp_foo_val = load ptr, ptr null, align 8
    ; CHECK: [[LABEL_TMP0:\.L.*]]:
-  call void (...) %fp_foo_val(), !callee_type !0
+  call void (...) %fp_foo_val(), !callee_type !0   
+  call void @direct_foo()
   %fp_bar_val = load ptr, ptr null, align 8
   ; CHECK: [[LABEL_TMP1:\.L.*]]:
-  %call_fp_bar = call i32 %fp_bar_val(i8 0), !callee_type !2
+  %call_fp_bar = call i32 %fp_bar_val(i8 0), !callee_type !2  
+  %call_fp_bar_direct = call i32 @direct_bar(i8 1)
   %fp_baz_val = load ptr, ptr null, align 8
   ; CHECK: [[LABEL_TMP2:\.L.*]]:
   %call_fp_baz = call ptr %fp_baz_val(ptr null), !callee_type !4
+  call void @direct_foo()
+  %call_fp_baz_direct = call ptr @direct_baz(ptr null)
+  call void @direct_foo()
   ret ptr %call_fp_baz
 }
 
@@ -41,3 +53,8 @@ entry:
 ;; Test for MD5 hash of _ZTSFPvS_E.generalized and the generated temporary callsite label.
 ; CHECK-NEXT: .quad   8646233951371320954
 ; CHECK-NEXT: .quad   [[LABEL_TMP2]]
+;; Test for number of direct calls and {callsite_label, callee} pairs.
+; CHECK-NEXT: .quad	3
+; CHECK-NEXT: .quad	direct_foo
+; CHECK-NEXT: .quad	direct_bar
+; CHECK-NEXT: .quad	direct_baz