[AArch64, compiler-rt] Implement trampoline intrinsics

[iamlouk]

Summary

The llvm.init.trampoline intrinsic is custom lowered to a call of
__trampoline_setup in the AArch64 instruction selection. As these
trampolines require an executable stack, a fatal error is reported
on Android or AArch64 Darwin platforms instead.

The function __trampoline_setup is implemented in the compiler-rt
for the AArch64 architecture.

This patch is inspired by the implementation of the same
intrinsics on the PowerPC architecture.

[compnerd]

compiler-rt/lib/builtins/trampoline_setup.c

45

Does this apply to Windows? I think that the condition here is incorrect.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5775

This will emit this on Windows as well, is this supposed to work there?

[iamlouk]

compiler-rt/lib/builtins/trampoline_setup.c

45

The x86 trampoline intrinsic (which inlines this logic) does work on Windows, but sadly I was not able to find out if Windows on ARM still allowes an executable stack via the online documentation I looked through, nor do I have the possibility to test it myself. I will search some more, but if I do not find a clear answer, I will change this condition here and the one where the call is emitted to exclude Windows in the next differential.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5775

Same as above where the trampoline setup function is implemented: I will change this to exclude windows in the next differential I upload.

[iamlouk]

Rebase and Windows on AArch64 treatment

The trampoline intrinsic is not treated for Windows on AArch64 anymore.

[compnerd]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5776

Is this intrinsic available in GCC? When was it added there? This could be an ABI break otherwise. We should check the runtime that we are going to link against before lowering this I think.

llvm/test/CodeGen/AArch64/trampoline.ll

35

Could you add negative checks for the other platforms please?

[iamlouk]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5776

No, GCC/libgcc does not have this for AArch64. GCC actually writes a trampoline template into the text section, then copies it over onto the stack of the function that calls GCCs equivalent of init.trampoline, and then creates instructions that store the address of the function and nest parameter into the copied-over template.

On x86 for example, inlining the trampoline creation is very simple because there is a instruction for loading a 64 bit immediate, but the lack thereof makes this tricky on AArch64.

I can give a version that is more similar to what GCC does a try, if that is preferred. That would remove the call to compiler-rt (and the function there), but make the lowering of the intrinsic a lot more complex.

It is maybe worth pointing out that clang does not actually implement nested functions, and will never create these intrinsics, so I don't know what ABI would break here.

(Sorry for the very late reply)

[compnerd]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

5776

That makes this even less valuable - GCC will not create the call, clang will not create the function. I think that doing what GCC does here is the correct thing to do.

[iamlouk]

Thank you for your feedback compnerd. Changing the implementation to mimic GCC's way of handling it's equivalent of this intrinsic requires a completely new approach though, so I am abandoning this revision.

[kiranchandramohan]

A new patch was put up on phabricator for adding trampolines to AArch64. Not sure whether it addresses the concerns raised here. @iamlouk
https://github.com/llvm/llvm-project/pull/70267

compiler-rt/lib/builtins/trampoline_setup.c

COMPILER_RT_ABI void __trampoline_setup(uint32_t *trampOnStack,
   trampOnStack[7] = 0x7d8903a6; // mtctr r12
   trampOnStack[8] = 0x816b0004; // lwz    r11,4(r11)
   trampOnStack[9] = 0x4e800420; // bctr
 
   // clear instruction cache
   __clear_cache(trampOnStack, &trampOnStack[10]);
 }
 #endif // __powerpc__ && !defined(__powerpc64__)
+
+#if defined(__aarch64__) && !defined(__ANDROID__) && !defined(__APPLE__) && \
+    !defined(_WIN64)
+static __inline void aarch64_gen_constant(uint32_t *buf, uint64_t constant,
+                                          uint64_t reg) {
+  // mov  reg, #<bits 15...0>
+  buf[0] = 0xd2800000u | (((constant >> 0) & 0xffffu) << 5) | reg;
+  // movk reg, #<bits 31...16>, lsl #16
+  buf[1] = 0xf2a00000u | (((constant >> 16) & 0xffffu) << 5) | reg;
+  // movk reg, #<bits 47...32>, lsl #32
+  buf[2] = 0xf2c00000u | (((constant >> 32) & 0xffffu) << 5) | reg;
+  // movk reg, #<bits 63...48>, lsl #48
+  buf[3] = 0xf2e00000u | (((constant >> 48) & 0xffffu) << 5) | reg;
+}
+
+COMPILER_RT_ABI void __trampoline_setup(uint32_t *trampOnStack,
+          uint64_t trampSizeAllocated,
+                                        const void *realFunc, void *localsPtr) {
+  if (trampSizeAllocated < 36) // 9 32-bit instructions.
+    compilerrt_abort();
+
+  // store realFunc in x9
+  aarch64_gen_constant(trampOnStack, (uint64_t)realFunc, 0x09);
+  // store localsPtr in x18
+  aarch64_gen_constant(&trampOnStack[4], (uint64_t)localsPtr, 0x12);
+  // br x9
+  trampOnStack[8] = 0xd61f0120;
+
+  // clear instruction cache
+  __clear_cache(trampOnStack, &trampOnStack[9]);
+}
+#endif /* defined(__aarch64__) && !defined(__ANDROID__) &&
+          !defined(__APPLE__) && !defined(_WIN64) */

compiler-rt/test/builtins/Unit/trampoline_setup_test.c

 // RUN: %clang_builtins %s %librt -fnested-functions -o %t && %run %t
 // REQUIRES: librt_has_trampoline_setup
 
 #include <stdio.h>
 #include <string.h>
 #include <stdint.h>
 
 /*
  * Tests nested functions
- * The ppc compiler generates a call to __trampoline_setup
+ * The ppc and aarch64 compiler generates a call to __trampoline_setup
  * The i386 and x86_64 compilers generate a call to ___enable_execute_stack
  */
 
 /*
  * Note that, nested functions are not ISO C and are not supported in Clang.
  */
 
 #if !defined(__clang__)

llvm/lib/Target/AArch64/AArch64ISelLowering.h

 //==-- AArch64ISelLowering.h - AArch64 DAG Lowering Interface ----*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
SDValue LowerFixedLengthConcatVectorsToSVE(SDValue Op,
                                              SelectionDAG &DAG) const;
   SDValue LowerFixedLengthFPExtendToSVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFixedLengthFPRoundToSVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFixedLengthIntToFPToSVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFixedLengthFPToIntToSVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFixedLengthVECTOR_SHUFFLEToSVE(SDValue Op,
                                               SelectionDAG &DAG) const;
 
+  SDValue LowerINIT_TRAMPOLINE(SDValue Op, SelectionDAG &DAG) const;
+
   SDValue BuildSDIVPow2(SDNode *N, const APInt &Divisor, SelectionDAG &DAG,
                         SmallVectorImpl<SDNode *> &Created) const override;
   SDValue BuildSREMPow2(SDNode *N, const APInt &Divisor, SelectionDAG &DAG,
                         SmallVectorImpl<SDNode *> &Created) const override;
   SDValue getSqrtEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                           int &ExtraSteps, bool &UseOneConst,
                           bool Reciprocal) const override;
   SDValue getRecipEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 //===-- AArch64ISelLowering.cpp - AArch64 DAG Lowering Implementation  ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
for (unsigned im = (unsigned)ISD::PRE_INC;
     setIndexedStoreAction(im, MVT::bf16, Legal);
   }
 
   // Trap.
   setOperationAction(ISD::TRAP, MVT::Other, Legal);
   setOperationAction(ISD::DEBUGTRAP, MVT::Other, Legal);
   setOperationAction(ISD::UBSANTRAP, MVT::Other, Legal);
 
+  // llvm.init.trampoline and llvm.adjust.trampoline are custom lowered.
+  setOperationAction(ISD::INIT_TRAMPOLINE, MVT::Other, Custom);
+  setOperationAction(ISD::ADJUST_TRAMPOLINE, MVT::Other, Custom);
+
   // We combine OR nodes for bitfield operations.
   setTargetDAGCombine(ISD::OR);
   // Try to create BICs for vector ANDs.
   setTargetDAGCombine(ISD::AND);
 
   // Vector add and sub nodes may conceal a high-half opportunity.
   // Also, try to fold ADD into CSINC/CSINV..
   setTargetDAGCombine({ISD::ADD, ISD::ABS, ISD::SUB, ISD::XOR, ISD::SINT_TO_FP,
if (SDValue Cmp = emitConjunction(DAG, Cond, CC)) {
     SDValue CCVal = DAG.getConstant(CC, dl, MVT::i32);
     return DAG.getNode(AArch64ISD::BRCOND, dl, MVT::Other, Chain, Dest, CCVal,
                        Cmp);
   }
 
   return SDValue();
 }
 
+// Lower @llvm.init.trampoline to a call to the function __trampoline_setup()
+SDValue AArch64TargetLowering::LowerINIT_TRAMPOLINE(SDValue Op,
+                                                    SelectionDAG &DAG) const {
+  // Because executable stacks are needed for trampolines:
+  if (Subtarget->isTargetAndroid() || Subtarget->isTargetDarwin() ||
+      Subtarget->isTargetWindows())
+    report_fatal_error("trampoline intrinsics are "
+                       "not supported on this platform.");
+
+  SDValue Root = Op.getOperand(0);
+  SDValue Trmp = Op.getOperand(1); // trampoline
+  SDValue FPtr = Op.getOperand(2); // nested function
+  SDValue Nest = Op.getOperand(3); // 'nest' parameter
+
+  SDLoc dl(Op);
+  EVT PtrVT = getPointerTy(DAG.getDataLayout());
+  TargetLowering::ArgListTy Args;
+
+  // Lower to a call to __trampoline_setup
+  TargetLowering::ArgListEntry Entry;
+  Entry.Ty = DAG.getDataLayout().getIntPtrType(*DAG.getContext());
+  Entry.Node = Trmp;
+  Args.push_back(Entry);
+  Entry.Node = DAG.getConstant(36, dl, MVT::i64);
+  Args.push_back(Entry);
+  Entry.Node = FPtr;
+  Args.push_back(Entry);
+  Entry.Node = Nest;
+  Args.push_back(Entry);
+
+  TargetLowering::CallLoweringInfo CLI(DAG);
+  CLI.setDebugLoc(dl).setChain(Root).setLibCallee(
+      CallingConv::C, Type::getVoidTy(*DAG.getContext()),
+      DAG.getExternalSymbol("__trampoline_setup", PtrVT), std::move(Args));
+
+  return LowerCallTo(CLI).second;
+}
+
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
 
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
std::pair<SDValue, SDValue> Pair =
         DAG.SplitScalar(Op.getOperand(2), DL, MVT::i64, MVT::i64);
 
     // chain = MSRR(chain, sysregname, lo, hi)
     SDValue Result = DAG.getNode(AArch64ISD::MSRR, DL, MVT::Other, Chain,
                                  SysRegName, Pair.first, Pair.second);
 
     return Result;
   }
+  case ISD::ADJUST_TRAMPOLINE:
+    return Op.getOperand(0);
+  case ISD::INIT_TRAMPOLINE:
+    return LowerINIT_TRAMPOLINE(Op, DAG);
   }
 }
 
 bool AArch64TargetLowering::mergeStoresAfterLegalization(EVT VT) const {
   return !Subtarget->useSVEForFixedLengthVectors();
 }
 
 bool AArch64TargetLowering::useSVEForFixedLengthVectorVT(

llvm/test/CodeGen/AArch64/trampoline.ll

+; RUN: llc -mtriple=aarch64-- < %s | FileCheck %s
+; UNSUPPORTED: android, darwin, system-windows
+
+declare void @llvm.init.trampoline(i8*, i8*, i8*);
+declare i8* @llvm.adjust.trampoline(i8*);
+
+define i32 @foo(i32* nest %ptr, i32 %val)
+{
+    %x = load i32, i32* %ptr
+    %sum = add i32 %x, %val
+    ret i32 %sum
+}
+
+; CHECK-LABEL: main
+define i32 @main(i32, i8**)
+{
+    %closure = alloca i32
+    store i32 13, i32* %closure
+    %closure_ptr = bitcast i32* %closure to i8*
+
+    %tramp_buf = alloca [36 x i8], align 4
+    %tramp_ptr = getelementptr [36 x i8], [36 x i8]* %tramp_buf, i32 0, i32 0
+; CHECK:  bl  __trampoline_setup
+    call void @llvm.init.trampoline(
+            i8* %tramp_ptr,
+            i8* bitcast (i32 (i32*, i32)* @foo to i8*),
+            i8* %closure_ptr)
+    %ptr = call i8* @llvm.adjust.trampoline(i8* %tramp_ptr)
+    %fp = bitcast i8* %ptr to i32(i32)*
+
+    %val2 = call i32 %fp (i32 42)
+
+    ret i32 %val2
+}
+