MachineFunction: Copy call site info when duplicating insts

Summary:
Preserve call site info for duplicated instructions. We copy over the
call site info in CloneMachineInstrBundle to avoid repeated calls to
copyCallSiteInfo in CloneMachineInstr.

(Alternatively, we could copy call site info higher up the stack, e.g.
into TargetInstrInfo::duplicate, or even into individual backend passes.
However, I don't see how that would be safer or more general than the
current approach.)

Reviewers: aprantl, djtodoro, dstenb

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D77685

llvm/lib/CodeGen/MachineFunction.cpp

@@ -426,6 +426,11 @@ MachineInstr &MachineFunction::CloneMachineInstrBundle(MachineBasicBlock &MBB,
       break;
     ++I;
   }
+  // Copy over call site info to the cloned instruction if needed. If Orig is in
+  // a bundle, copyCallSiteInfo takes care of finding the call instruction in
+  // the bundle.
+  if (Orig.shouldUpdateCallSiteInfo())
+    copyCallSiteInfo(&Orig, FirstClone);
   return *FirstClone;
 }
 

llvm/test/CodeGen/X86/taildup-callsiteinfo.mir

@@ -0,0 +1,75 @@
+# RUN: llc %s -emit-call-site-info -run-pass=block-placement -tail-dup-placement-threshold=4 -o - | FileCheck %s
+#
+# Test case adapted from test/CodeGen/X86/taildup-heapallocsite.ll.
+
+# CHECK-LABEL: callSites:
+# CHECK-NEXT:   - { bb: 1, offset: 1, fwdArgRegs:
+# CHECK-NEXT:       - { arg: 0, reg: '$rcx' } }
+# CHECK-NEXT:   - { bb: 2, offset: 1, fwdArgRegs:
+# CHECK-NEXT:       - { arg: 0, reg: '$rcx' } }
+
+--- |
+  target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+  target triple = "x86_64-pc-windows-msvc19.22.27905"
+
+  define dso_local void @taildupit(i32* readonly %size_ptr) {
+  entry:
+    %tobool = icmp eq i32* %size_ptr, null
+    br i1 %tobool, label %cond.end, label %cond.true
+
+  cond.true:                                        ; preds = %entry
+    %0 = load i32, i32* %size_ptr, align 4
+    br label %cond.end
+
+  cond.end:                                         ; preds = %cond.true, %entry
+    %cond = phi i32 [ %0, %cond.true ], [ 1, %entry ]
+    %call = tail call i8* @alloc(i32 %cond)
+    tail call void @f2()
+    ret void
+  }
+
+  declare dso_local i8* @alloc(i32)
+
+  declare dso_local void @f2()
+
+...
+---
+name:            taildupit
+tracksRegLiveness: true
+liveins:
+  - { reg: '$rcx', virtual-reg: '' }
+callSites:
+  - { bb: 3, offset: 0, fwdArgRegs:
+      - { arg: 0, reg: '$rcx' } }
+body:             |
+  bb.0.entry:
+    successors: %bb.1(0x30000000), %bb.2(0x50000000)
+    liveins: $rcx
+
+    $rsp = frame-setup SUB64ri8 $rsp, 40, implicit-def dead $eflags
+    frame-setup SEH_StackAlloc 40
+    frame-setup SEH_EndPrologue
+    TEST64rr renamable $rcx, renamable $rcx, implicit-def $eflags
+    JCC_1 %bb.2, 5, implicit killed $eflags
+
+  bb.1:
+    successors: %bb.3(0x80000000)
+
+    renamable $ecx = MOV32ri 1
+    JMP_1 %bb.3
+
+  bb.2.cond.true:
+    successors: %bb.3(0x80000000)
+    liveins: $rcx
+
+    renamable $ecx = MOV32rm killed renamable $rcx, 1, $noreg, 0, $noreg :: (load 4 from %ir.size_ptr)
+
+  bb.3.cond.end:
+    liveins: $ecx
+
+    CALL64pcrel32 @alloc, csr_win64, implicit $rsp, implicit $ssp, implicit $ecx, implicit-def $rsp, implicit-def $ssp, implicit-def dead $rax
+    SEH_Epilogue
+    $rsp = frame-destroy ADD64ri8 $rsp, 40, implicit-def dead $eflags
+    TAILJMPd64 @f2, csr_win64, implicit $rsp, implicit $ssp, implicit $rsp, implicit $ssp
+
+...

llvm/unittests/CodeGen/MFCommon.inc

@@ -92,11 +92,18 @@ private:
   TargetInstrInfo TII;
 };
 
+static TargetOptions getTargetOptionsForBogusMachine() {
+  TargetOptions Opts;
+  Opts.EmitCallSiteInfo = true;
+  return Opts;
+}
+
 class BogusTargetMachine : public LLVMTargetMachine {
 public:
   BogusTargetMachine()
-      : LLVMTargetMachine(Target(), "", Triple(""), "", "", TargetOptions(),
-                          Reloc::Static, CodeModel::Small, CodeGenOpt::Default),
+      : LLVMTargetMachine(Target(), "", Triple(""), "", "",
+                          getTargetOptionsForBogusMachine(), Reloc::Static,
+                          CodeModel::Small, CodeGenOpt::Default),
         ST(*this) {}
 
   ~BogusTargetMachine() override {}

llvm/unittests/CodeGen/MachineInstrTest.cpp

@@ -383,6 +383,26 @@ TEST(MachineInstrExtraInfo, RemoveExtraInfo) {
   ASSERT_FALSE(MI->getHeapAllocMarker());
 }
 
+TEST(MachineInstrClone, CopyCallSiteInfo) {
+  LLVMContext Ctx;
+  Module Mod("Module", Ctx);
+  auto MF = createMachineFunction(Ctx, Mod);
+  auto MBB = MF->CreateMachineBasicBlock();
+  auto MII = MBB->begin();
+
+  MCInstrDesc MCID = {0, 0,       0,       0,      0, (1ULL << MCID::Call),
+                      0, nullptr, nullptr, nullptr};
+
+  MachineFunction::CallSiteInfo CSInfo;
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  ASSERT_TRUE(MI->isCandidateForCallSiteEntry());
+  MBB->insert(MII, MI);
+  MF->addCallArgsForwardingRegs(MI, std::move(CSInfo));
+  EXPECT_EQ(MF->getCallSitesInfo().size(), 1u);
+  MF->CloneMachineInstrBundle(*MBB, MBB->end(), *MI);
+  EXPECT_EQ(MF->getCallSitesInfo().size(), 2u);
+}
+
 static_assert(is_trivially_copyable<MCOperand>::value, "trivially copyable");
 
 } // end namespace