[AArch64] Fix the size passed to __trampoline_setup

[ssijaric-nv]

The trampoline size is 36 bytes on AArch64. The runtime function __trampoline_setup
aborts as it expects the trampoline size of at least 36 bytes, and the size passed
is 20 bytes. Fix the inconsistency in AArch64TargetLowering::LowerINIT_TRAMPOLINE.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: None (ssijaric-nv)

Changes

The trampoline size is 36 bytes on AArch64. The runtime function __trampoline_setup
aborts as it expects the trampoline size of at least 36 bytes, and the size passed
is 20 bytes. Fix the inconsistency in AArch64TargetLowering::LowerINIT_TRAMPOLINE.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+9-2)
• (modified) llvm/test/CodeGen/AArch64/trampoline.ll (+1)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index e1be825fcf7bf3..418cef075ba65f 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7148,9 +7148,16 @@ SDValue AArch64TargetLowering::LowerINIT_TRAMPOLINE(SDValue Op,
   Entry.Ty = IntPtrTy;
   Entry.Node = Trmp;
   Args.push_back(Entry);
-  Entry.Node = DAG.getConstant(20, dl, MVT::i64);
-  Args.push_back(Entry);
 
+  if (auto *FI = dyn_cast<FrameIndexSDNode>(Trmp.getNode())) {
+    MachineFunction &MF = DAG.getMachineFunction();
+    MachineFrameInfo &MFI = MF.getFrameInfo();
+    Entry.Node =
+        DAG.getConstant(MFI.getObjectSize(FI->getIndex()), dl, MVT::i64);
+  } else
+    Entry.Node = DAG.getConstant(36, dl, MVT::i64);
+
+  Args.push_back(Entry);
   Entry.Node = FPtr;
   Args.push_back(Entry);
   Entry.Node = Nest;
diff --git a/llvm/test/CodeGen/AArch64/trampoline.ll b/llvm/test/CodeGen/AArch64/trampoline.ll
index 293e538a7459d4..217cda1a4ce1ca 100644
--- a/llvm/test/CodeGen/AArch64/trampoline.ll
+++ b/llvm/test/CodeGen/AArch64/trampoline.ll
@@ -12,6 +12,7 @@ define i64 @main() {
   %val = alloca i64
   %nval = bitcast ptr %val to ptr
   %tramp = alloca [36 x i8], align 8
+  ; CHECK:	mov	w1, #36
   ; CHECK:	bl	__trampoline_setup
   call void @llvm.init.trampoline(ptr %tramp, ptr @f, ptr %nval)
   %fp = call ptr @llvm.adjust.trampoline(ptr %tramp)

[Leporacanthicus]

LGTM.

[davemgreen]

This looks like it matches https://github.com/llvm/llvm-project/blob/main/compiler-rt/lib/builtins/trampoline_setup.c.

Is it worth adding a test for the FrameIndex size and static versions?

[ssijaric-nv]

Thanks, David. I have only seen cases where we go through the FrameIndex path as the trampoline is on the stack. I can add a test where the trampoline is a global belonging to a section that can be executed, which will trigger the 'else' check. Or perhaps turn the 'else' into an assert?

[davemgreen]

Thanks for adding the new test. Still LGTM.