[llvm] [CodeGen][X86] Fix lowering of tailcalls when `-ms-hotpatch` is used (PR #77245)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

Author: Alexandre Ganea (aganea)

<details>
<summary>Changes</summary>

Previously, tail jump pseudo-opcodes were skipped by the `encodeInstruction()` call inside `X86AsmPrinter::LowerPATCHABLE_OP`. This caused emission of a 2-byte NOP and dropping of the tail jump.

Also, the `PatchableFunction` pass didn't properly update/erase the call site info for the existing `MachineInstr`, after wrapping it into a `TargetOpcode::PATCHABLE_OP`.

---
Full diff: https://github.com/llvm/llvm-project/pull/77245.diff


3 Files Affected:

- (modified) llvm/lib/CodeGen/PatchableFunction.cpp (+3) 
- (modified) llvm/lib/Target/X86/X86MCInstLower.cpp (+6-2) 
- (added) llvm/test/CodeGen/X86/patchable-prologue-tailcall.ll (+69) 


``````````diff
diff --git a/llvm/lib/CodeGen/PatchableFunction.cpp b/llvm/lib/CodeGen/PatchableFunction.cpp
index 9449f143366f0f..768e8dff2efc24 100644
--- a/llvm/lib/CodeGen/PatchableFunction.cpp
+++ b/llvm/lib/CodeGen/PatchableFunction.cpp
@@ -87,6 +87,9 @@ bool PatchableFunction::runOnMachineFunction(MachineFunction &MF) {
   for (auto &MO : FirstActualI->operands())
     MIB.add(MO);
 
+  if (FirstActualI->shouldUpdateCallSiteInfo())
+    MF.eraseCallSiteInfo(&*FirstActualI);
+
   FirstActualI->eraseFromParent();
   MF.ensureAlignment(Align(16));
   return true;
diff --git a/llvm/lib/Target/X86/X86MCInstLower.cpp b/llvm/lib/Target/X86/X86MCInstLower.cpp
index e1a67f61e76640..93aa0fc23eb9f2 100644
--- a/llvm/lib/Target/X86/X86MCInstLower.cpp
+++ b/llvm/lib/Target/X86/X86MCInstLower.cpp
@@ -959,7 +959,8 @@ void X86AsmPrinter::LowerPATCHABLE_OP(const MachineInstr &MI,
   bool EmptyInst = (Opcode == TargetOpcode::PATCHABLE_OP);
 
   MCInst MCI;
-  MCI.setOpcode(Opcode);
+  // Make sure below we don't encode pseudo tailcalls.
+  MCI.setOpcode(convertTailJumpOpcode(Opcode));
   for (auto &MO : drop_begin(MI.operands(), 2))
     if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, MO))
       MCI.addOperand(*MaybeOperand);
@@ -994,8 +995,11 @@ void X86AsmPrinter::LowerPATCHABLE_OP(const MachineInstr &MI,
       (void)NopSize;
     }
   }
-  if (!EmptyInst)
+  if (!EmptyInst) {
+    if (Opcode != convertTailJumpOpcode(Opcode))
+      OutStreamer->AddComment("TAILCALL");
     OutStreamer->emitInstruction(MCI, getSubtargetInfo());
+  }
 }
 
 // Lower a stackmap of the form:
diff --git a/llvm/test/CodeGen/X86/patchable-prologue-tailcall.ll b/llvm/test/CodeGen/X86/patchable-prologue-tailcall.ll
new file mode 100644
index 00000000000000..ab3ff5852a74b0
--- /dev/null
+++ b/llvm/test/CodeGen/X86/patchable-prologue-tailcall.ll
@@ -0,0 +1,69 @@
+; RUN: llc -verify-machineinstrs < %s | FileCheck %s --check-prefix=CHECK
+
+; CHECK: "?mi_new_test@@YAPEAX_K at Z":
+; CHECK-NEXT: # %bb.0:
+; CHECK-NEXT: jmp     mi_new
+
+; CHECK: "?builtin_malloc_test@@YAPEAX_K at Z":
+; CHECK-NEXT: # %bb.0:
+; CHECK-NEXT: jmp     malloc
+
+; // Built with: clang-cl.exe /c patchable-prologue-tailcall.cpp /O2 /hotpatch -Xclang -emit-llvm
+;
+; typedef unsigned long long size_t;
+; 
+; extern "C" {
+;     void* mi_new(size_t size);
+;  }
+;
+; void *mi_new_test(size_t count)
+; {
+;     return mi_new(count);
+; }
+; 
+; void *builtin_malloc_test(size_t count)
+; {
+;     return __builtin_malloc(count);
+; }
+
+; ModuleID = 'patchable-prologue-tailcall.cpp'
+source_filename = "patchable-prologue-tailcall.cpp"
+target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.38.33133"
+
+; Function Attrs: mustprogress nounwind sspstrong uwtable
+define dso_local noundef ptr @"?mi_new_test@@YAPEAX_K at Z"(i64 noundef %count) local_unnamed_addr #0 {
+entry:
+  %call = tail call ptr @mi_new(i64 noundef %count) #4
+  ret ptr %call
+}
+
+declare dso_local ptr @mi_new(i64 noundef) local_unnamed_addr #1
+
+; Function Attrs: mustprogress nofree nounwind sspstrong willreturn memory(inaccessiblemem: readwrite) uwtable
+define dso_local noalias noundef ptr @"?builtin_malloc_test@@YAPEAX_K at Z"(i64 noundef %count) local_unnamed_addr #2 {
+entry:
+  %call = tail call ptr @malloc(i64 noundef %count) #4
+  ret ptr %call
+}
+
+; Function Attrs: mustprogress nofree nounwind willreturn allockind("alloc,uninitialized") allocsize(0) memory(inaccessiblemem: readwrite)
+declare dso_local noalias noundef ptr @malloc(i64 noundef) local_unnamed_addr #3
+
+attributes #0 = { mustprogress nounwind sspstrong uwtable "min-legal-vector-width"="0" "no-trapping-math"="true" "patchable-function"="prologue-short-redirect" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #1 = { "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #2 = { mustprogress nofree nounwind sspstrong willreturn memory(inaccessiblemem: readwrite) uwtable "min-legal-vector-width"="0" "no-trapping-math"="true" "patchable-function"="prologue-short-redirect" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #3 = { mustprogress nofree nounwind willreturn allockind("alloc,uninitialized") allocsize(0) memory(inaccessiblemem: readwrite) "alloc-family"="malloc" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #4 = { nounwind }
+
+!llvm.linker.options = !{!0, !1}
+!llvm.module.flags = !{!2, !3, !4, !5}
+!llvm.ident = !{!6}
+
+!0 = !{!"/DEFAULTLIB:libcmt.lib"}
+!1 = !{!"/DEFAULTLIB:oldnames.lib"}
+!2 = !{i32 1, !"wchar_size", i32 2}
+!3 = !{i32 8, !"PIC Level", i32 2}
+!4 = !{i32 7, !"uwtable", i32 2}
+!5 = !{i32 1, !"MaxTLSAlign", i32 65536}
+!6 = !{!"clang version 18.0.0git (https://github.com/llvm/llvm-project.git 0ccf9e29e9626a3a6813b35608c8959178c50a6f)"}

``````````

</details>


https://github.com/llvm/llvm-project/pull/77245