[CallSiteInfo][CallGraphSection] Extend CallSiteInfo with TypeId

[Prabhuk]

Summary

Add TypeId field to CallSiteInfo for tracking callee type id. Read the type
id in and out by the MIR parser/printer.

With this patch, TypeId is only set from callSites.typeId fields from MIR
inputs. An upcoming patch will pass type ids from the clang front-end,
and they will be read and set while lowering in the LLVM middle-end.

Original RFC: https://lists.llvm.org/pipermail/llvm-dev/2021-June/151044.html
Updated RFC: https://lists.llvm.org/pipermail/llvm-dev/2021-July/151739.html

[morehouse]

LGTM

[vitalybuka]

llvm/include/llvm/CodeGen/MIRYamlMapping.h

460

Why this is not in the operator==

same for ArgForwardingRegs

[arsenm]

Should add some parse failure condition tests

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

406

Don't need hasValue

408

• instead of getValue

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

7

Should check the actual values are preserved

11–32

Does this really need the IR section?

[arsenm]

Reverse ping

[Prabhuk]

https://discourse.llvm.org/t/rfc-computing-storing-and-restoring-conservative-call-graphs-with-llvm/58446/3?u=prabhuk

[Prabhuk]

Rebased the patchset and addressed the compilation failures

llvm/include/llvm/CodeGen/MIRYamlMapping.h

struct MachineInstrLoc {
     bool operator==(const MachineInstrLoc &Other) const {
       return BlockNum == Other.BlockNum && Offset == Other.Offset;
     }
   };
 
   MachineInstrLoc CallLocation;
   std::vector<ArgRegPair> ArgForwardingRegs;
 
+  /// Numeric callee type identifier used for call graph section.
+  using TypeIdTy = Optional<uint64_t>;
+  TypeIdTy TypeId;
+
   bool operator==(const CallSiteInfo &Other) const {
     return CallLocation.BlockNum == Other.CallLocation.BlockNum &&
            CallLocation.Offset == Other.CallLocation.Offset;
   }
 };
 
 template <> struct MappingTraits<CallSiteInfo::ArgRegPair> {
   static void mapping(IO &YamlIO, CallSiteInfo::ArgRegPair &ArgReg) {
     YamlIO.mapRequired("arg", ArgReg.ArgNo);

 namespace yaml {
 
 template <> struct MappingTraits<CallSiteInfo> {
   static void mapping(IO &YamlIO, CallSiteInfo &CSInfo) {
     YamlIO.mapRequired("bb", CSInfo.CallLocation.BlockNum);
     YamlIO.mapRequired("offset", CSInfo.CallLocation.Offset);
     YamlIO.mapOptional("fwdArgRegs", CSInfo.ArgForwardingRegs,
                        std::vector<CallSiteInfo::ArgRegPair>());
+    YamlIO.mapOptional("typeId", CSInfo.TypeId, None);
   }
 
   static const bool flow = true;
 };
 
 /// Serializable representation of debug value substitutions.
 struct DebugValueSubstitution {
   unsigned SrcInst;

llvm/include/llvm/CodeGen/MachineFunction.h

struct ArgRegPair {
     ArgRegPair(Register R, unsigned Arg) : Reg(R), ArgNo(Arg) {
       assert(Arg < (1 << 16) && "Arg out of range");
     }
   };
 
   struct CallSiteInfo {
     /// Vector of call argument and its forwarding register.
     SmallVector<ArgRegPair, 1> ArgRegPairs;
+
+    /// Callee type id.
+    ConstantInt *TypeId = nullptr;
   };
 
 private:
   Delegate *TheDelegate = nullptr;
   GISelChangeObserver *Observer = nullptr;
 
   using CallSiteInfoMap = DenseMap<const MachineInstr *, CallSiteInfo>;
-  /// Map a call instruction to call site arguments forwarding info.
+  /// Map a call instruction to call site arguments forwarding and type id.
   CallSiteInfoMap CallSitesInfo;
 
   /// A helper function that returns call site info for a give call
   /// instruction if debug entry value support is enabled.
   CallSiteInfoMap::iterator getCallSiteInfo(const MachineInstr *MI);
 
   // Callbacks for insertion and removal.
   void handleInsertion(MachineInstr &MI);
void setVariableDbgInfo(const DILocalVariable *Var, const DIExpression *Expr,
     VariableDbgInfos.emplace_back(Var, Expr, Slot, Loc);
   }
 
   VariableDbgInfoMapTy &getVariableDbgInfo() { return VariableDbgInfos; }
   const VariableDbgInfoMapTy &getVariableDbgInfo() const {
     return VariableDbgInfos;
   }
 
-  /// Start tracking the arguments passed to the call \p CallI.
+  /// Start tracking the arguments passed to the call \p CallI and call type.
   void addCallSiteInfo(const MachineInstr *CallI, CallSiteInfo &&CallInfo) {
     assert(CallI->isCandidateForCallSiteEntry());
     bool Inserted =
         CallSitesInfo.try_emplace(CallI, std::move(CallInfo)).second;
     (void)Inserted;
     assert(Inserted && "Call site info not unique");
   }
 

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

if (!CallI->isCall(MachineInstr::IgnoreBundle))
                    " is not a call instruction");
     MachineFunction::CallSiteInfo CSInfo;
     for (auto ArgRegPair : YamlCSInfo.ArgForwardingRegs) {
       Register Reg;
       if (parseNamedRegisterReference(PFS, Reg, ArgRegPair.Reg.Value, Error))
         return error(Error, ArgRegPair.Reg.SourceRange);
       CSInfo.ArgRegPairs.emplace_back(Reg, ArgRegPair.ArgNo);
     }
+    if (YamlCSInfo.TypeId.hasValue()) {
+      IntegerType *Int64Ty = Type::getInt64Ty(Context);
+      CSInfo.TypeId = ConstantInt::get(Int64Ty, YamlCSInfo.TypeId.getValue(),
+                                       /*isSigned=*/false);
+    }
 
-    if (TM.Options.EmitCallSiteInfo)
+    if (TM.Options.EmitCallSiteInfo || TM.Options.EmitCallGraphSection)
       MF.addCallSiteInfo(&*CallI, std::move(CSInfo));
   }
 
-  if (YamlMF.CallSitesInfo.size() && !TM.Options.EmitCallSiteInfo)
+  if (YamlMF.CallSitesInfo.size() &&
+      !(TM.Options.EmitCallSiteInfo || TM.Options.EmitCallGraphSection))
     return error(Twine("Call site info provided but not used"));
   return false;
 }
 
 void MIRParserImpl::setupDebugValueTracking(
     MachineFunction &MF, PerFunctionMIParsingState &PFS,
     const yaml::MachineFunction &YamlMF) {
   // Compute the value of the "next instruction number" field.

llvm/lib/CodeGen/MIRPrinter.cpp

for (auto CSInfo : MF.getCallSitesInfo()) {
     YmlCS.CallLocation = CallLocation;
     // Construct call arguments and theirs forwarding register info.
     for (auto ArgReg : CSInfo.second.ArgRegPairs) {
       yaml::CallSiteInfo::ArgRegPair YmlArgReg;
       YmlArgReg.ArgNo = ArgReg.ArgNo;
       printRegMIR(ArgReg.Reg, YmlArgReg.Reg, TRI);
       YmlCS.ArgForwardingRegs.emplace_back(YmlArgReg);
     }
+    // Get type id.
+    if (CSInfo.second.TypeId)
+      YmlCS.TypeId = CSInfo.second.TypeId->getZExtValue();
     YMF.CallSitesInfo.push_back(YmlCS);
   }
 
   // Sort call info by position of call instructions.
   llvm::sort(YMF.CallSitesInfo.begin(), YMF.CallSitesInfo.end(),
              [](yaml::CallSiteInfo A, yaml::CallSiteInfo B) {
                if (A.CallLocation.BlockNum == B.CallLocation.BlockNum)
                  return A.CallLocation.Offset < B.CallLocation.Offset;

llvm/lib/CodeGen/MachineFunction.cpp

try_next:;
   return FilterID;
 }
 
 MachineFunction::CallSiteInfoMap::iterator
 MachineFunction::getCallSiteInfo(const MachineInstr *MI) {
   assert(MI->isCandidateForCallSiteEntry() &&
          "Call site info refers only to call (MI) candidates");
 
-  if (!Target.Options.EmitCallSiteInfo)
+  if (!Target.Options.EmitCallSiteInfo && !Target.Options.EmitCallGraphSection)
     return CallSitesInfo.end();
   return CallSitesInfo.find(MI);
 }
 
 /// Return the call machine instruction or find a call within bundle.
 static const MachineInstr *getCallInstr(const MachineInstr *MI) {
   if (!MI->isBundle())
     return MI;

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

+# Test MIR printer and parser for type id field in callSites. It is used
+# for propogating call site type identifiers to emit in the call graph section.
+
+# RUN: llc --call-graph-section %s -run-pass=none -o - | FileCheck %s
+# CHECK: name: main
+# CHECK: callSites:
+# CHECK-NEXT: - { bb: {{.*}}, offset: {{.*}}, fwdArgRegs: [], typeId:
+# CHECK-NEXT: 123456789 }
+
+--- |
+  ; ModuleID = 'test.ll'
+  source_filename = "test.ll"
+  target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+  target triple = "x86_64-unknown-linux-gnu"
+  
+  define dso_local void @foo(i8 signext %a) {
+  entry:
+    ret void
+  }
+  
+  define dso_local i32 @main() {
+  entry:
+    %retval = alloca i32, align 4
+    %fp = alloca void (i8)*, align 8
+    store i32 0, i32* %retval, align 4
+    store void (i8)* @foo, void (i8)** %fp, align 8
+    %0 = load void (i8)*, void (i8)** %fp, align 8
+    call void %0(i8 signext 97)
+    ret i32 0
+  }
+
+...
+---
+name:            foo
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    RET 0
+
+...
+---
+name:            main
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: retval, type: default, offset: 0, size: 4, alignment: 4, 
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true, 
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+  - { id: 1, name: fp, type: default, offset: 0, size: 8, alignment: 8, 
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true, 
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+callSites:
+  - { bb: 0, offset: 6, fwdArgRegs: [], typeId: 
+    123456789 }
+body:             |
+  bb.0.entry:
+    MOV32mi %stack.0.retval, 1, $noreg, 0, $noreg, 0 :: (store (s32) into %ir.retval)
+    MOV64mi32 %stack.1.fp, 1, $noreg, 0, $noreg, @foo :: (store (s64) into %ir.fp)
+    %0:gr64 = MOV32ri64 @foo
+    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    %1:gr32 = MOV32ri 97
+    $edi = COPY %1
+    CALL64r killed %0, csr_64, implicit $rsp, implicit $ssp, implicit $edi, implicit-def $rsp, implicit-def $ssp
+    ADJCALLSTACKUP64 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    %2:gr32 = MOV32r0 implicit-def dead $eflags
+    $eax = COPY %2
+    RET 0, $eax
+
+...