[llvm][AArch64] Fix BTI after returns_twice when call has no attributes

[DavidSpickett]

Summary

Previously we were looking for the returns twice attribute by manually
getting the function attributes from the call. This meant that we only
found attributes on the call itself, not what it was calling.

So if you had:
%call1 = call i32 @setjmp(ptr noundef null)

We would not BTI protect that even though setjmp clearly needs it.

Clang happens to produce:
%call = call i32 @setjmp(ptr noundef null) #0 ; returns_twice

So all valid calls were protected. This is not guaranteed,
the frontend may choose not to put attributes on the call.

It is undefined behaviour to call setjmp indirectly
(https://pubs.opengroup.org/onlinepubs/9699919799/functions/setjmp.html)
but as I misused the APIs here I think it's worth fixing up regardless.

Added comments to the test file where the IR differs from what
clang would output.

[DavidSpickett]

As suggested on https://github.com/llvm/llvm-project/issues/60732.

With this fix if clang converts an indirect call into a direct call, we will correctly protect that with a bti instruction. This matches the expected behaviour and what gcc produces.

[nikic]

LG

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   bool IsThisReturn = false;
 
   AArch64FunctionInfo *FuncInfo = MF.getInfo<AArch64FunctionInfo>();
   bool TailCallOpt = MF.getTarget().Options.GuaranteedTailCallOpt;
   bool IsCFICall = CLI.CB && CLI.CB->isIndirectCall() && CLI.CFIType;
   bool IsSibCall = false;
   bool GuardWithBTI = false;
 
-  if (CLI.CB && CLI.CB->getAttributes().hasFnAttr(Attribute::ReturnsTwice) &&
+  if (CLI.CB && CLI.CB->hasFnAttr(Attribute::ReturnsTwice) &&
       !Subtarget->noBTIAtReturnTwice()) {
     GuardWithBTI = FuncInfo->branchTargetEnforcement();
   }
 
   // Analyze operands of the call, assigning locations to each operand.
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(CallConv, IsVarArg, MF, ArgLocs, *DAG.getContext());
 

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   unsigned Opc = 0;
   // Calls with operand bundle "clang.arc.attachedcall" are special. They should
   // be expanded to the call, directly followed by a special marker sequence and
   // a call to an ObjC library function.
   if (Info.CB && objcarc::hasAttachedCallOpBundle(Info.CB))
     Opc = AArch64::BLR_RVMARKER;
   // A call to a returns twice function like setjmp must be followed by a bti
   // instruction.
-  else if (Info.CB &&
-           Info.CB->getAttributes().hasFnAttr(Attribute::ReturnsTwice) &&
+  else if (Info.CB && Info.CB->hasFnAttr(Attribute::ReturnsTwice) &&
            !Subtarget.noBTIAtReturnTwice() &&
            MF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
     Opc = AArch64::BLR_BTI;
   else
     Opc = getCallOpcode(MF, Info.Callee.isReg(), false);
 
   auto MIB = MIRBuilder.buildInstrNoInsert(Opc);
   unsigned CalleeOpNo = 0;

llvm/test/CodeGen/AArch64/setjmp-bti.ll

 
 ; C source
 ; --------
 ; extern int setjmp(ptr);
 ; extern void notsetjmp(void);
 ;
 ; void bbb(void) {
 ;   setjmp(0);
+;   setjmp(0); // With the attributes removed.
 ;   int (*fnptr)(ptr) = setjmp;
-;   fnptr(0);
+;   fnptr(0); // With attributes added.
 ;   notsetjmp();
 ; }
 
 define void @bbb() {
 ; BTI-LABEL: bbb:
 ; BTI:       bl setjmp
 ; BTI-NEXT:  hint #36
+; BTI:       bl setjmp
+; BTI-NEXT:  hint #36
 ; BTI:       blr x{{[0-9]+}}
 ; BTI-NEXT:  hint #36
 ; BTI:       bl notsetjmp
 ; BTI-NOT:   hint #36
 
 ; BTISLS-LABEL: bbb:
 ; BTISLS:       bl setjmp
 ; BTISLS-NEXT:  hint #36
+; BTISLS:       bl setjmp
+; BTISLS-NEXT:  hint #36
 ; BTISLS:       bl __llvm_slsblr_thunk_x{{[0-9]+}}
 ; BTISLS-NEXT:  hint #36
 ; BTISLS:       bl notsetjmp
 ; BTISLS-NOT:   hint #36
 
 ; NOBTI-LABEL: bbb:
 ; NOBTI:     bl setjmp
 ; NOBTI-NOT: hint #36
+; NOBTI:     bl setjmp
+; NOBTI-NOT: hint #36
 ; NOBTI:     blr x{{[0-9]+}}
 ; NOBTI-NOT: hint #36
 ; NOBTI:     bl notsetjmp
 ; NOBTI-NOT: hint #36
 entry:
   %fnptr = alloca ptr, align 8
+  ; The frontend may apply attributes to the call, but it doesn't have to. We
+  ; should be looking at the call base, which looks past that to the called function.
   %call = call i32 @setjmp(ptr noundef null) #0
+  %call1 = call i32 @setjmp(ptr noundef null)
   store ptr @setjmp, ptr %fnptr, align 8
   %0 = load ptr, ptr %fnptr, align 8
-  %call1 = call i32 %0(ptr noundef null) #0
+  ; Clang does not attach the attribute here but if it did, it should work.
+  %call2 = call i32 %0(ptr noundef null) #0
   call void @notsetjmp()
   ret void
 }
 
 declare i32 @setjmp(ptr noundef) #0
 declare void @notsetjmp()
 
 attributes #0 = { returns_twice }
 
 !llvm.module.flags = !{!0}
 !0 = !{i32 8, !"branch-target-enforcement", i32 1}