llvm · Prabhuk · Apr 3, 2024 · Apr 22, 2024 · Apr 24, 2024 · Apr 29, 2024
diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -182,6 +182,36 @@ class AsmPrinter : public MachineFunctionPass {
  /// Emit comments in assembly output if this is true.
  bool VerboseAsm;

+  /// Store symbols and type identifiers used to create call graph section
+  /// entries related to a function.
+  struct FunctionInfo {
+    /// Numeric type identifier used in call graph section for indirect calls
+    /// and targets.
+    using CGTypeId = uint64_t;
+
+    /// Enumeration of function kinds, and their mapping to function kind values
+    /// stored in call graph 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 only for indirect
+    /// calls and inclusive of all indirect calls of the function.
+    SmallVector<std::pair<CGTypeId, MCSymbol *>> CallSiteLabels;
+  };
+
+  enum CallGraphSectionFormatVersion : uint64_t {
+    V_0 = 0,
+  };
+
  /// Output stream for the stack usage file (i.e., .su file).
  std::unique_ptr<raw_fd_ostream> StackUsageStream;

@@ -342,6 +372,13 @@ class 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,
+      const MachineFunction::CallSiteInfoMap &CallSitesInfoMap,
+      const MachineInstr &MI);
+
  //===------------------------------------------------------------------===//
  // XRay instrumentation implementation.
  //===------------------------------------------------------------------===//
@@ -429,6 +466,8 @@ class 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 emitPseudoProbe(const MachineInstr &MI);

  void emitRemarksSection(remarks::RemarkStreamer &RS);

diff --git a/llvm/include/llvm/CodeGen/MachineFunction.h b/llvm/include/llvm/CodeGen/MachineFunction.h
@@ -534,11 +534,12 @@ class LLVM_ABI MachineFunction {
    unsigned TargetFlags;
  };

+  using CallSiteInfoMap = DenseMap<const MachineInstr *, CallSiteInfo>;
+
 private:
  Delegate *TheDelegate = nullptr;
  GISelChangeObserver *Observer = nullptr;

-  using CallSiteInfoMap = DenseMap<const MachineInstr *, CallSiteInfo>;
  /// Map a call instruction to call site arguments forwarding info.
  CallSiteInfoMap CallSitesInfo;


diff --git a/llvm/include/llvm/MC/MCObjectFileInfo.h b/llvm/include/llvm/MC/MCObjectFileInfo.h
@@ -68,6 +68,9 @@ class MCObjectFileInfo {
  /// Language Specific Data Area information is emitted to.
  MCSection *LSDASection = nullptr;

+  /// Section containing call graph metadata.
+  MCSection *CallGraphSection = nullptr;
+
  /// If exception handling is supported by the target and the target can
  /// support a compact representation of the CIE and FDE, this is the section
  /// to emit them into.
@@ -360,6 +363,8 @@ class MCObjectFileInfo {
  MCSection *getFaultMapSection() const { return FaultMapSection; }
  MCSection *getRemarksSection() const { return RemarksSection; }

+  MCSection *getCallGraphSection(const MCSection &TextSec) const;
+
  MCSection *getStackSizesSection(const MCSection &TextSec) const;

  MCSection *getBBAddrMapSection(const MCSection &TextSec) const;

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1640,6 +1640,94 @@ void AsmPrinter::emitStackUsage(const MachineFunction &MF) {
    *StackUsageStream << "static\n";
 }

+/// Extracts a generalized numeric type identifier of a Function's type from
+/// type metadata. Returns null if metadata cannot be found.
+static ConstantInt *extractNumericCGTypeId(const Function &F) {
+  SmallVector<MDNode *, 2> Types;
+  F.getMetadata(LLVMContext::MD_type, Types);
+  for (const auto &Type : Types) {
+    if (Type->hasGeneralizedMDString()) {
+      MDString *MDGeneralizedTypeId = cast<MDString>(Type->getOperand(1));
+      uint64_t TypeIdVal = llvm::MD5Hash(MDGeneralizedTypeId->getString());
+      IntegerType *Int64Ty = Type::getInt64Ty(F.getContext());
+      return ConstantInt::get(Int64Ty, TypeIdVal);
+    }
+  }
+
+  return nullptr;
+}
+
+/// Emits .callgraph section.
+void AsmPrinter::emitCallGraphSection(const MachineFunction &MF,
+                                      FunctionInfo &FuncInfo) {
+  if (!MF.getTarget().Options.EmitCallGraphSection)
+    return;
+
+  // Switch to the call graph section for the function
+  MCSection *FuncCGSection =
+      getObjFileLowering().getCallGraphSection(*getCurrentSection());
+  assert(FuncCGSection && "null callgraph section");
+  OutStreamer->pushSection();
+  OutStreamer->switchSection(FuncCGSection);
+
+  // Emit format version number.
+  OutStreamer->emitInt64(CallGraphSectionFormatVersion::V_0);
+
+  // Emit function's self information, which is composed of:
+  //  1) FunctionEntryPc
+  //  2) FunctionKind: Whether the function is indirect target, and if so,
+  //     whether its type id is known.
+  //  3) FunctionTypeId: Emit only when the function is an indirect target
+  //     and its type id is known.
+
+  // Emit function entry pc.
+  const MCSymbol *FunctionSymbol = getFunctionBegin();
+  OutStreamer->emitSymbolValue(FunctionSymbol, TM.getProgramPointerSize());
+
+  // If this function has external linkage or has its address taken and
+  // it is not a callback, then anything could call it.
+  const Function &F = MF.getFunction();
+  bool IsIndirectTarget =
+      !F.hasLocalLinkage() || F.hasAddressTaken(nullptr,
+                                                /*IgnoreCallbackUses=*/true,
+                                                /*IgnoreAssumeLikeCalls=*/true,
+                                                /*IgnoreLLVMUsed=*/false);
+
+  // FIXME: FunctionKind takes a few values but emitted as a 64-bit value.
+  // Can be optimized to occupy 2 bits instead.
+  // Emit function kind, and type id if available.
+  if (!IsIndirectTarget) {
+    OutStreamer->emitInt64(
+        static_cast<uint64_t>(FunctionInfo::FunctionKind::NOT_INDIRECT_TARGET));
+  } else {
+    const auto *TypeId = extractNumericCGTypeId(F);
+    if (TypeId) {
+      OutStreamer->emitInt64(static_cast<uint64_t>(
+          FunctionInfo::FunctionKind::INDIRECT_TARGET_KNOWN_TID));
+      OutStreamer->emitInt64(TypeId->getZExtValue());
+    } else {
+      OutStreamer->emitInt64(static_cast<uint64_t>(
+          FunctionInfo::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;
+
+  // Emit the count of pairs.
+  OutStreamer->emitInt64(CallSiteLabels.size());
+
+  // Emit the type id and call site label pairs.
+  for (const auto &[TypeId, Label] : CallSiteLabels) {
+    OutStreamer->emitInt64(TypeId);
+    OutStreamer->emitSymbolValue(Label, TM.getProgramPointerSize());
+  }
+  FuncInfo.CallSiteLabels.clear();
+
+  OutStreamer->popSection();
+}
+
 void AsmPrinter::emitPCSectionsLabel(const MachineFunction &MF,
                                     const MDNode &MD) {
  MCSymbol *S = MF.getContext().createTempSymbol("pcsection");
@@ -1770,6 +1858,28 @@ static StringRef getMIMnemonic(const MachineInstr &MI, MCStreamer &Streamer) {
  return Name;
 }

+void AsmPrinter::emitIndirectCalleeLabels(
+    FunctionInfo &FuncInfo,
+    const MachineFunction::CallSiteInfoMap &CallSitesInfoMap,
+    const MachineInstr &MI) {
+  // Only indirect calls have type identifiers set.
+  const auto &CallSiteInfo = CallSitesInfoMap.find(&MI);
+  if (CallSiteInfo == CallSitesInfoMap.end())
+    return;
+
+  for (ConstantInt *CalleeTypeId : CallSiteInfo->second.CalleeTypeIds) {
+    // Emit label.
+    MCSymbol *S = MF->getContext().createTempSymbol();
+    OutStreamer->emitLabel(S);
+
+    // Get numeric callee_type id value.
+    uint64_t CalleeTypeIdVal = CalleeTypeId->getZExtValue();
+
+    // Add to function's callsite labels.
+    FuncInfo.CallSiteLabels.emplace_back(CalleeTypeIdVal, S);
+  }
+}
+
 /// EmitFunctionBody - This method emits the body and trailer for a
 /// function.
 void AsmPrinter::emitFunctionBody() {
@@ -1816,6 +1926,8 @@ void AsmPrinter::emitFunctionBody() {
    MBBSectionRanges[MF->front().getSectionID()] =
        MBBSectionRange{CurrentFnBegin, nullptr};

+  FunctionInfo FuncInfo;
+  const auto &CallSitesInfoMap = MF->getCallSitesInfo();
  for (auto &MBB : *MF) {
    // Print a label for the basic block.
    emitBasicBlockStart(MBB);
@@ -1946,6 +2058,9 @@ void AsmPrinter::emitFunctionBody() {
        break;
      }

+      if (TM.Options.EmitCallGraphSection && MI.isCall())
+        emitIndirectCalleeLabels(FuncInfo, CallSitesInfoMap, MI);
+
      // If there is a post-instruction symbol, emit a label for it here.
      if (MCSymbol *S = MI.getPostInstrSymbol())
        OutStreamer->emitLabel(S);
@@ -2125,6 +2240,9 @@ void AsmPrinter::emitFunctionBody() {
  // Emit section containing stack size metadata.
  emitStackSizeSection(*MF);

+  // Emit section containing call graph metadata.
+  emitCallGraphSection(*MF, FuncInfo);
+
  // Emit .su file containing function stack size information.
  emitStackUsage(*MF);

@@ -2799,6 +2917,7 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
      F.hasFnAttribute("xray-instruction-threshold") ||
      needFuncLabels(MF, *this) || NeedsLocalForSize ||
      MF.getTarget().Options.EmitStackSizeSection ||
+      MF.getTarget().Options.EmitCallGraphSection ||
      MF.getTarget().Options.BBAddrMap) {
    CurrentFnBegin = createTempSymbol("func_begin");
    if (NeedsLocalForSize)

diff --git a/llvm/lib/MC/MCObjectFileInfo.cpp b/llvm/lib/MC/MCObjectFileInfo.cpp
@@ -536,6 +536,8 @@ void MCObjectFileInfo::initELFMCObjectFileInfo(const Triple &T, bool Large) {
  EHFrameSection =
      Ctx->getELFSection(".eh_frame", EHSectionType, EHSectionFlags);

+  CallGraphSection = Ctx->getELFSection(".callgraph", ELF::SHT_PROGBITS, 0);
+
  StackSizesSection = Ctx->getELFSection(".stack_sizes", ELF::SHT_PROGBITS, 0);

  PseudoProbeSection = Ctx->getELFSection(".pseudo_probe", DebugSecType, 0);
@@ -1081,6 +1083,24 @@ MCSection *MCObjectFileInfo::getDwarfComdatSection(const char *Name,
  llvm_unreachable("Unknown ObjectFormatType");
 }

+MCSection *
+MCObjectFileInfo::getCallGraphSection(const MCSection &TextSec) const {
+  if (Ctx->getObjectFileType() != MCContext::IsELF)
+    return CallGraphSection;
+
+  const MCSectionELF &ElfSec = static_cast<const MCSectionELF &>(TextSec);
+  unsigned Flags = ELF::SHF_LINK_ORDER;
+  StringRef GroupName;
+  if (const MCSymbol *Group = ElfSec.getGroup()) {
+    GroupName = Group->getName();
+    Flags |= ELF::SHF_GROUP;
+  }
+
+  return Ctx->getELFSection(".callgraph", ELF::SHT_PROGBITS, Flags, 0,
+                            GroupName, true, ElfSec.getUniqueID(),
+                            cast<MCSymbolELF>(TextSec.getBeginSymbol()));
+}
+
 MCSection *
 MCObjectFileInfo::getStackSizesSection(const MCSection &TextSec) const {
  if ((Ctx->getObjectFileType() != MCContext::IsELF) ||

diff --git a/llvm/test/CodeGen/X86/call-graph-section-assembly.ll b/llvm/test/CodeGen/X86/call-graph-section-assembly.ll
@@ -0,0 +1,60 @@
+;; Test if temporary labels are generated for each indirect callsite with a callee_type metadata.
+;; Test if the .callgraph section contains the numerical callee type id for each of the temporary 
+;; labels generated. 
+
+; RUN: llc -mtriple=x86_64-unknown-linux --call-graph-section -o - < %s | FileCheck %s
+
+declare !type !0 void @foo()
+
+declare !type !1 i32 @bar(i8 signext)
+
+declare !type !2 ptr @baz(ptr)
+
+; CHECK: ball:
+; CHECK-NEXT: .Lfunc_begin0:
+define void @ball() {
+entry:
+  %retval = alloca i32, align 4
+  %fp_foo = alloca ptr, align 8
+  %a = alloca i8, align 1
+  %fp_bar = alloca ptr, align 8
+  %fp_baz = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp_foo, align 8
+  %fp_foo_val = load ptr, ptr %fp_foo, align 8
+  ; CHECK: .Ltmp0:
+  call void (...) %fp_foo_val(), !callee_type !1
+  store ptr @bar, ptr %fp_bar, align 8
+  %fp_bar_val = load ptr, ptr %fp_bar, align 8
+  %a_val = load i8, ptr %a, align 1
+  ; CHECK: .Ltmp1:
+  %call_fp_bar = call i32 %fp_bar_val(i8 signext %a_val), !callee_type !3
+  store ptr @baz, ptr %fp_baz, align 8
+  %fp_baz_val = load ptr, ptr %fp_baz, align 8
+  ; CHECK: .Ltmp2:
+  %call_fp_baz = call ptr %fp_baz_val(ptr %a), !callee_type !5
+  call void @foo()
+  %a_val_2 = load i8, ptr %a, align 1
+  %call_bar = call i32 @bar(i8 signext %a_val_2)
+  %call_baz = call ptr @baz(ptr %a)
+  ret void
+}
+
+; CHECK: .section .callgraph,"o",@progbits,.text
+
+; CHECK-NEXT: .quad   0
+; CHECK-NEXT: .quad   .Lfunc_begin0
+; CHECK-NEXT: .quad   1
+; CHECK-NEXT: .quad   3
+; CHECK-NEXT: .quad   4524972987496481828
+; CHECK-NEXT: .quad   .Ltmp0
+!0 = !{i64 0, !"_ZTSFvE.generalized"}
+!1 = !{!0}
+; CHECK-NEXT: .quad   3498816979441845844
+; CHECK-NEXT: .quad   .Ltmp1
+!2 = !{i64 0, !"_ZTSFicE.generalized"}
+!3 = !{!2}
+; CHECK-NEXT: .quad   8646233951371320954
+; CHECK-NEXT: .quad   .Ltmp2
+!4 = !{i64 0, !"_ZTSFPvS_E.generalized"}
+!5 = !{!4}
diff --git a/llvm/test/CodeGen/X86/call-graph-section-tailcall.ll b/llvm/test/CodeGen/X86/call-graph-section-tailcall.ll
@@ -0,0 +1,34 @@
+;; Tests that we store the type identifiers in .callgraph section of the object file for tailcalls.
+
+; RUN: llc -mtriple=x86_64-unknown-linux --call-graph-section -filetype=obj -o - < %s | \
+; RUN: llvm-readelf -x .callgraph - | FileCheck %s
+
+define i32 @_Z13call_indirectPFicEc(ptr %func, i8 %x) !type !0 {
+entry:
+  %call = tail call i32 %func(i8 signext %x), !callee_type !1
+  ret i32 %call
+}
+
+define i32 @main(i32 %argc) !type !3 {
+entry:
+  %0 = and i32 %argc, 1
+  %cmp = icmp eq i32 %0, 0
+  %_Z3fooc._Z3barc = select i1 %cmp, ptr @_Z3fooc, ptr @_Z3barc
+  %call.i = tail call i32 %_Z3fooc._Z3barc(i8 signext 97), !callee_type !1
+  ret i32 %call.i
+}
+
+declare !type !2 i32 @_Z3fooc(i8 signext)
+
+declare !type !2 i32 @_Z3barc(i8 signext)
+
+;; Check that the numeric type id (md5 hash) for the below type ids are emitted
+;; to the callgraph section.
+
+; CHECK: Hex dump of section '.callgraph':
+
+!0 = !{i64 0, !"_ZTSFiPvcE.generalized"}
+!1 = !{!2}
+; CHECK-DAG: 5486bc59 814b8e30
+!2 = !{i64 0, !"_ZTSFicE.generalized"}
+!3 = !{i64 0, !"_ZTSFiiE.generalized"}