[AArch64][GlobalISel] Fix llvm.returnaddress(0) selection when LR is clobbered.

aemerson · aemerson · commit 2e25d75aaa22 · 2020-01-21T22:53:32.000-08:00
The code was originally ported from SelectionDAG, which does CSE behind the scenes
automatically. When copying the return address from LR live into the function, we
need to make sure to use the single copy on function entry. Any later copy from LR
could be using clobbered junk.

Implement this by caching the copy in the per-MF state in the selector.

Should hopefully fix the AArch64 sanitiser buildbot failure.
diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -63,6 +63,7 @@ class AArch64InstructionSelector : public InstructionSelector {
     // cache it here for each run of the selector.
     ProduceNonFlagSettingCondBr =
         !MF.getFunction().hasFnAttribute(Attribute::SpeculativeLoadHardening);
+    MFReturnAddr = Register();
   }
 
 private:
@@ -123,7 +124,7 @@ class AArch64InstructionSelector : public InstructionSelector {
                                 MachineRegisterInfo &MRI) const;
   bool selectIntrinsicWithSideEffects(MachineInstr &I,
                                       MachineRegisterInfo &MRI) const;
-  bool selectIntrinsic(MachineInstr &I, MachineRegisterInfo &MRI) const;
+  bool selectIntrinsic(MachineInstr &I, MachineRegisterInfo &MRI);
   bool selectVectorICmp(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectIntrinsicTrunc(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectIntrinsicRound(MachineInstr &I, MachineRegisterInfo &MRI) const;
@@ -295,6 +296,11 @@ class AArch64InstructionSelector : public InstructionSelector {
 
   bool ProduceNonFlagSettingCondBr = false;
 
+  // Some cached values used during selection.
+  // We use LR as a live-in register, and we keep track of it here as it can be
+  // clobbered by calls.
+  Register MFReturnAddr;
+
 #define GET_GLOBALISEL_PREDICATES_DECL
 #include "AArch64GenGlobalISel.inc"
 #undef GET_GLOBALISEL_PREDICATES_DECL
@@ -4079,8 +4085,8 @@ bool AArch64InstructionSelector::selectIntrinsicWithSideEffects(
   return true;
 }
 
-bool AArch64InstructionSelector::selectIntrinsic(
-    MachineInstr &I, MachineRegisterInfo &MRI) const {
+bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
+                                                 MachineRegisterInfo &MRI) {
   unsigned IntrinID = findIntrinsicID(I);
   if (!IntrinID)
     return false;
@@ -4138,10 +4144,20 @@ bool AArch64InstructionSelector::selectIntrinsic(
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
 
     if (Depth == 0 && IntrinID == Intrinsic::returnaddress) {
+      if (MFReturnAddr) {
+        MIRBuilder.buildCopy({DstReg}, MFReturnAddr);
+        I.eraseFromParent();
+        return true;
+      }
       MFI.setReturnAddressIsTaken(true);
       MF.addLiveIn(AArch64::LR, &AArch64::GPR64spRegClass);
       I.getParent()->addLiveIn(AArch64::LR);
-      MIRBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+      // Insert the copy from LR/X30 into the entry block, before it can be
+      // clobbered by anything.
+      MachineIRBuilder EntryBuilder(MF);
+      EntryBuilder.setInstr(*MF.begin()->begin());
+      EntryBuilder.buildCopy({DstReg}, {Register(AArch64::LR)});
+      MFReturnAddr = DstReg;
       I.eraseFromParent();
       return true;
     }
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-returnaddr.ll b/llvm/test/CodeGen/AArch64/GlobalISel/select-returnaddr.ll
@@ -9,6 +9,20 @@ entry:
   ret i8* %0
 }
 
+define i8* @rt0_call_clobber(i32 %x) nounwind readnone {
+entry:
+; CHECK-LABEL: rt0_call_clobber:
+; CHECK: stp     x20, x19, [sp, #-32]!
+; CHECK: stp     x29, x30, [sp, #16]
+; CHECK: mov     x19, x30
+; CHECK: bl      _foo
+; CHECK: ldp     x29, x30, [sp, #16]
+; CHECK: mov     x0, x19
+  %ret = call i32 @foo()
+  %0 = tail call i8* @llvm.returnaddress(i32 0)
+  ret i8* %0
+}
+
 define i8* @rt2() nounwind readnone {
 entry:
 ; CHECK-LABEL: rt2:
@@ -19,4 +33,6 @@ entry:
   ret i8* %0
 }
 
+
+declare i32 @foo()
 declare i8* @llvm.returnaddress(i32) nounwind readnone
