[llvm] Extract and propagate indirect call type id

Extract numeric type ids for indirect calls, and carry them to the back-end
with CallSiteInfo. The numeric type ids will be used at the back-end to emit
the call graph section.

Pull Request: llvm#87575

Author: necipfazil

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -26,11 +26,14 @@
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/EHPersonalities.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/ArrayRecycler.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/MD5.h"
 #include "llvm/Support/Recycler.h"
 #include "llvm/Target/TargetOptions.h"
 #include <bitset>
@@ -499,6 +502,37 @@ class LLVM_ABI MachineFunction {
 
     /// Callee type id.
     ConstantInt *TypeId = nullptr;
+
+    CallSiteInfo() = default;
+
+    /// Extracts the numeric type id from the CallBase's type operand bundle,
+    /// and sets TypeId. This is used as type id for the indirect call in the
+    /// call graph section.
+    CallSiteInfo(const CallBase &CB) {
+      // Call graph section needs numeric type id only for indirect calls.
+      if (!CB.isIndirectCall())
+        return;
+
+      std::optional<OperandBundleUse> Opt =
+          CB.getOperandBundle(LLVMContext::OB_callee_type);
+      // Return if the operand bundle for call graph section cannot be found.
+      if (!Opt)
+        return;
+
+      // Get generalized type id string
+      auto OB = *Opt;
+      assert(OB.Inputs.size() == 1 && "invalid input size");
+      auto *OBVal = OB.Inputs.front().get();
+      auto *TypeIdMD = cast<MetadataAsValue>(OBVal)->getMetadata();
+      auto *TypeIdStr = cast<MDString>(TypeIdMD);
+      assert(TypeIdStr->getString().ends_with(".generalized") &&
+             "invalid type identifier");
+
+      // Compute numeric type id from generalized type id string
+      uint64_t TypeIdVal = MD5Hash(TypeIdStr->getString());
+      IntegerType *Int64Ty = Type::getInt64Ty(CB.getContext());
+      TypeId = ConstantInt::get(Int64Ty, TypeIdVal, /*IsSigned=*/false);
+    }
   };
 
   struct CalledGlobalInfo {

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

@@ -889,7 +889,8 @@ EmitSchedule(MachineBasicBlock::iterator &InsertPos) {
     }
 
     if (MI->isCandidateForAdditionalCallInfo()) {
-      if (DAG->getTarget().Options.EmitCallSiteInfo)
+      if (DAG->getTarget().Options.EmitCallSiteInfo ||
+          DAG->getTarget().Options.EmitCallGraphSection)
         MF.addCallSiteInfo(MI, DAG->getCallSiteInfo(Node));
 
       if (auto CalledGlobal = DAG->getCalledGlobal(Node))

llvm/lib/IR/Verifier.cpp

@@ -3794,6 +3794,11 @@ void Verifier::visitCallBase(CallBase &Call) {
     } else if (Tag == LLVMContext::OB_callee_type) {
       Check(!FoundCalleeTypeBundle, "Multiple \"callee_type\" operand bundles",
             Call);
+      auto *OBVal = BU.Inputs.front().get();
+      auto *TypeIdMD = cast<MetadataAsValue>(OBVal)->getMetadata();
+      auto *TypeIdStr = cast<MDString>(TypeIdMD);
+      Check(TypeIdStr->getString().ends_with(".generalized"),
+            "Invalid \"callee_type\" type identifier", Call);
       FoundCalleeTypeBundle = true;
     }
   }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8915,6 +8915,7 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   bool &IsTailCall = CLI.IsTailCall;
   CallingConv::ID &CallConv = CLI.CallConv;
   bool IsVarArg = CLI.IsVarArg;
+  const auto *CB = CLI.CB;
 
   MachineFunction &MF = DAG.getMachineFunction();
   MachineFunction::CallSiteInfo CSInfo;
@@ -8954,6 +8955,10 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
                     *DAG.getContext());
   RetCCInfo.AnalyzeCallResult(Ins, RetCC);
 
+  // Set type id for call site info.
+  if (MF.getTarget().Options.EmitCallGraphSection && CB && CB->isIndirectCall())
+    CSInfo = MachineFunction::CallSiteInfo(*CB);
+
   // Check callee args/returns for SVE registers and set calling convention
   // accordingly.
   if (CallConv == CallingConv::C || CallConv == CallingConv::Fast) {

llvm/lib/Target/ARM/ARMISelLowering.cpp

@@ -2439,6 +2439,7 @@ ARMTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   CallingConv::ID CallConv              = CLI.CallConv;
   bool doesNotRet                       = CLI.DoesNotReturn;
   bool isVarArg                         = CLI.IsVarArg;
+  const auto *CB = CLI.CB;
 
   MachineFunction &MF = DAG.getMachineFunction();
   ARMFunctionInfo *AFI = MF.getInfo<ARMFunctionInfo>();
@@ -2462,6 +2463,10 @@ ARMTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
       !Subtarget->noBTIAtReturnTwice())
     GuardWithBTI = AFI->branchTargetEnforcement();
 
+  // Set type id for call site info.
+  if (MF.getTarget().Options.EmitCallGraphSection && CB && CB->isIndirectCall())
+    CSInfo = MachineFunction::CallSiteInfo(*CB);
+
   // Determine whether this is a non-secure function call.
   if (CLI.CB && CLI.CB->getAttributes().hasFnAttr("cmse_nonsecure_call"))
     isCmseNSCall = true;

llvm/lib/Target/Mips/MipsISelLowering.cpp

@@ -3264,6 +3264,7 @@ MipsTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   bool &IsTailCall                      = CLI.IsTailCall;
   CallingConv::ID CallConv              = CLI.CallConv;
   bool IsVarArg                         = CLI.IsVarArg;
+  const auto *CB = CLI.CB;
 
   MachineFunction &MF = DAG.getMachineFunction();
   MachineFrameInfo &MFI = MF.getFrameInfo();
@@ -3320,8 +3321,11 @@ MipsTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   // Get a count of how many bytes are to be pushed on the stack.
   unsigned StackSize = CCInfo.getStackSize();
 
-  // Call site info for function parameters tracking.
+  // Call site info for function parameters tracking and call base type info.
   MachineFunction::CallSiteInfo CSInfo;
+  // Set type id for call site info.
+  if (MF.getTarget().Options.EmitCallGraphSection && CB && CB->isIndirectCall())
+    CSInfo = MachineFunction::CallSiteInfo(*CB);
 
   // Check if it's really possible to do a tail call. Restrict it to functions
   // that are part of this compilation unit.

llvm/lib/Target/X86/X86FastISel.cpp

@@ -3631,6 +3631,12 @@ bool X86FastISel::fastLowerCall(CallLoweringInfo &CLI) {
   CLI.NumResultRegs = RVLocs.size();
   CLI.Call = MIB;
 
+  // Add call site info for call graph section.
+  if (TM.Options.EmitCallGraphSection && CB && CB->isIndirectCall()) {
+    MachineFunction::CallSiteInfo CSInfo(*CB);
+    MF->addCallSiteInfo(CLI.Call, std::move(CSInfo));
+  }
+
   return true;
 }
 
@@ -4026,6 +4032,8 @@ bool X86FastISel::tryToFoldLoadIntoMI(MachineInstr *MI, unsigned OpNo,
     MO.setReg(IndexReg);
   }
 
+  if (MI->isCall())
+    FuncInfo.MF->moveAdditionalCallInfo(MI, Result);
   Result->addMemOperand(*FuncInfo.MF, createMachineMemOperandFor(LI));
   Result->cloneInstrSymbols(*FuncInfo.MF, *MI);
   MachineBasicBlock::iterator I(MI);

llvm/lib/Target/X86/X86ISelLoweringCall.cpp

@@ -2031,6 +2031,10 @@ X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   if (CallConv == CallingConv::X86_INTR)
     report_fatal_error("X86 interrupts may not be called directly");
 
+  // Set type id for call site info.
+  if (MF.getTarget().Options.EmitCallGraphSection && CB && CB->isIndirectCall())
+    CSInfo = MachineFunction::CallSiteInfo(*CB);
+
   // Analyze operands of the call, assigning locations to each operand.
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(CallConv, isVarArg, MF, ArgLocs, *DAG.getContext());

llvm/test/CodeGen/AArch64/call-site-info-typeid.ll

@@ -0,0 +1,29 @@
+;; Tests that call site type ids can be extracted and set from type operand
+;; bundles.
+
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section -mtriple aarch64-linux-gnu < %s -stop-before=finalize-isel -o - | FileCheck %s
+
+define dso_local void @foo(i8 signext %a) !type !3 {
+entry:
+  ret void
+}
+
+; CHECK: name: main
+define dso_local i32 @main() !type !4 {
+entry:
+  %retval = alloca i32, align 4
+  %fp = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp, align 8
+  %0 = load ptr, ptr %fp, align 8
+  ; CHECK: callSites:
+  ; CHECK-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+  ; CHECK-NEXT: 7854600665770582568 }
+  call void %0(i8 signext 97) [ "callee_type"(metadata !"_ZTSFvcE.generalized") ]
+  ret i32 0
+}
+
+!3 = !{i64 0, !"_ZTSFvcE.generalized"}
+!4 = !{i64 0, !"_ZTSFiE.generalized"}

llvm/test/CodeGen/ARM/call-site-info-typeid.ll

@@ -0,0 +1,29 @@
+;; Tests that call site type ids can be extracted and set from type operand
+;; bundles.
+
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section -mtriple arm-linux-gnu < %s -stop-before=finalize-isel -o - | FileCheck %s
+
+define dso_local void @foo(i8 signext %a) !type !3 {
+entry:
+  ret void
+}
+
+; CHECK: name: main
+define dso_local i32 @main() !type !4 {
+entry:
+  %retval = alloca i32, align 4
+  %fp = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp, align 8
+  %0 = load ptr, ptr %fp, align 8
+  ; CHECK: callSites:
+  ; CHECK-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+  ; CHECK-NEXT: 7854600665770582568 }
+  call void %0(i8 signext 97) [ "callee_type"(metadata !"_ZTSFvcE.generalized") ]
+  ret i32 0
+}
+
+!3 = !{i64 0, !"_ZTSFvcE.generalized"}
+!4 = !{i64 0, !"_ZTSFiE.generalized"}

llvm/test/CodeGen/MIR/X86/call-site-info-typeid.ll

@@ -0,0 +1,102 @@
+;; Test MIR printer and parser for type id field in call site info. Test that
+;; it works well with/without --emit-call-site-info.
+
+;; Multiplex --call-graph-section and -emit-call-site-info as both utilize
+;; CallSiteInfo and callSites.
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test printer and parser with --call-graph-section only.
+
+;; Test printer.
+;; Verify that fwdArgRegs is not set, typeId is set.
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section %s -stop-before=finalize-isel -o %t1.mir
+; RUN: cat %t1.mir | FileCheck %s --check-prefix=PRINTER_CGS
+; PRINTER_CGS: name: main
+; PRINTER_CGS: callSites:
+; PRINTER_CGS-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+; PRINTER_CGS-NEXT: 7854600665770582568 }
+
+
+;; Test parser.
+;; Verify that we get the same result.
+; RUN: llc --call-graph-section %t1.mir -run-pass=finalize-isel -o - \
+; RUN: | FileCheck %s --check-prefix=PARSER_CGS
+; PARSER_CGS: name: main
+; PARSER_CGS: callSites:
+; PARSER_CGS-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+; PARSER_CGS-NEXT: 7854600665770582568 }
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test printer and parser with -emit-call-site-info only.
+
+;; Test printer.
+;; Verify that fwdArgRegs is set, typeId is not set.
+; RUN: llc -emit-call-site-info %s -stop-before=finalize-isel -o %t2.mir
+; RUN: cat %t2.mir | FileCheck %s --check-prefix=PRINTER_CSI
+; PRINTER_CSI: name: main
+; PRINTER_CSI: callSites:
+; PRINTER_CSI-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs:
+; PRINTER_CSI-NEXT: { arg: 0, reg: '$edi' }
+; PRINTER_CSI-NOT: typeId:
+
+
+;; Test parser.
+;; Verify that we get the same result.
+; RUN: llc -emit-call-site-info %t2.mir -run-pass=finalize-isel -o - \
+; RUN: | FileCheck %s --check-prefix=PARSER_CSI
+; PARSER_CSI: name: main
+; PARSER_CSI: callSites:
+; PARSER_CSI-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs:
+; PARSER_CSI-NEXT: { arg: 0, reg: '$edi' }
+; PARSER_CSI-NOT: typeId:
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Test printer and parser with both -emit-call-site-info and --call-graph-section.
+
+;; Test printer.
+;; Verify both fwdArgRegs and typeId are set.
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section -emit-call-site-info %s -stop-before=finalize-isel -o %t2.mir
+; RUN: cat %t2.mir | FileCheck %s --check-prefix=PRINTER_CGS_CSI
+; PRINTER_CGS_CSI: name: main
+; PRINTER_CGS_CSI: callSites:
+; PRINTER_CGS_CSI-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs:
+; PRINTER_CGS_CSI-NEXT: { arg: 0, reg: '$edi' }, typeId:
+; PRINTER_CGS_CSI-NEXT:   7854600665770582568 }
+
+
+;; Test parser.
+;; Verify that we get the same result.
+; RUN: llc --call-graph-section -emit-call-site-info %t2.mir -run-pass=finalize-isel -o - \
+; RUN: | FileCheck %s --check-prefix=PARSER_CGS_CSI
+; PARSER_CGS_CSI: name: main
+; PARSER_CGS_CSI: callSites:
+; PARSER_CGS_CSI-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs:
+; PARSER_CGS_CSI-NEXT: { arg: 0, reg: '$edi' }, typeId:
+; PARSER_CGS_CSI-NEXT:   7854600665770582568 }
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local void @foo(i8 signext %a) !type !3 {
+entry:
+  ret void
+}
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local i32 @main() !type !4 {
+entry:
+  %retval = alloca i32, align 4
+  %fp = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp, align 8
+  %0 = load ptr, ptr %fp, align 8
+  call void %0(i8 signext 97) [ "callee_type"(metadata !"_ZTSFvcE.generalized") ]
+  ret i32 0
+}
+
+!3 = !{i64 0, !"_ZTSFvcE.generalized"}
+!4 = !{i64 0, !"_ZTSFiE.generalized"}

llvm/test/CodeGen/Mips/call-site-info-typeid.ll

@@ -0,0 +1,29 @@
+;; Tests that call site type ids can be extracted and set from type operand
+;; bundles.
+
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section -mtriple=mips-linux-gnu < %s -stop-before=finalize-isel -o - | FileCheck %s
+
+define dso_local void @foo(i8 signext %a) !type !3 {
+entry:
+  ret void
+}
+
+; CHECK: name: main
+define dso_local i32 @main() !type !4 {
+entry:
+  %retval = alloca i32, align 4
+  %fp = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp, align 8
+  %0 = load ptr, ptr %fp, align 8
+  ; CHECK: callSites:
+  ; CHECK-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+  ; CHECK-NEXT: 7854600665770582568 }
+  call void %0(i8 signext 97) [ "callee_type"(metadata !"_ZTSFvcE.generalized") ]
+  ret i32 0
+}
+
+!3 = !{i64 0, !"_ZTSFvcE.generalized"}
+!4 = !{i64 0, !"_ZTSFiE.generalized"}

llvm/test/CodeGen/X86/call-site-info-typeid.ll

@@ -0,0 +1,29 @@
+;; Tests that call site type ids can be extracted and set from type operand
+;; bundles.
+
+;; Verify the exact typeId value to ensure it is not garbage but the value
+;; computed as the type id from the type operand bundle.
+; RUN: llc --call-graph-section -mtriple=x86_64-unknown-linux < %s -stop-before=finalize-isel -o - | FileCheck %s
+
+define dso_local void @foo(i8 signext %a) !type !3 {
+entry:
+  ret void
+}
+
+; CHECK: name: main
+define dso_local i32 @main() !type !4 {
+entry:
+  %retval = alloca i32, align 4
+  %fp = alloca ptr, align 8
+  store i32 0, ptr %retval, align 4
+  store ptr @foo, ptr %fp, align 8
+  %0 = load ptr, ptr %fp, align 8
+  ; CHECK: callSites:
+  ; CHECK-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+  ; CHECK-NEXT: 7854600665770582568 }
+  call void %0(i8 signext 97) [ "callee_type"(metadata !"_ZTSFvcE.generalized") ]
+  ret i32 0
+}
+
+!3 = !{i64 0, !"_ZTSFvcE.generalized"}
+!4 = !{i64 0, !"_ZTSFiE.generalized"}