[StackProtector] Allow targets to specify an instruction is part of terminator sequence

[nemanjai]

Summary

When the basic block is being split up to insert stack protector checks/traps, it is spliced prior to the instructions needed by the terminator. This includes implicit defs and register copies required by the ABI. However, on certain targets, there may be instructions that are part of such a terminator sequence but are not strictly speaking register copies.
Such is the case with the PowerPC "Move to count register" instruction. The instruction implicitly defines the count register that is used by the subsequent branch.

On PowerPC, if the target is not allowed to specify that the instruction is part of the terminator sequence, the block gets spliced between the instruction that sets up the CTR and the instruction that uses it (such as an indirect tail call).

[aemerson]

Can the test be reduced? Do you really need a main()?

[nemanjai]

In D125846#3521120, @aemerson wrote:

Can the test be reduced? Do you really need a main()?

Good point. In the original test case I reduced func1() is static so it needed the caller. I should be able to change the linkage of the function and remove main().

[arsenm]

Why can't you just mark these instructions as being terminators? AMDGPU has to glue a lot of ALU instructions to terminators, and we just define terminator and nonterminator pseudo variants

[nemanjai]

In D125846#3523688, @arsenm wrote:

Why can't you just mark these instructions as being terminators? AMDGPU has to glue a lot of ALU instructions to terminators, and we just define terminator and nonterminator pseudo variants

This seems like a less desirable path. Marking what is inherently a non-terminator instruction as a terminator seems likely to cause problems. I don't have a specific case in mind, but perhaps if some pass makes the assumption that it is safe to split a basic block into two that terminate with two terminators. It might be a safe thing to do except if some terminator is not actually a terminator.

[nemanjai]

In D125846#3523981, @nemanjai wrote:

In D125846#3523688, @arsenm wrote:

Why can't you just mark these instructions as being terminators? AMDGPU has to glue a lot of ALU instructions to terminators, and we just define terminator and nonterminator pseudo variants

This seems like a less desirable path. Marking what is inherently a non-terminator instruction as a terminator seems likely to cause problems. I don't have a specific case in mind, but perhaps if some pass makes the assumption that it is safe to split a basic block into two that terminate with two terminators. It might be a safe thing to do except if some terminator is not actually a terminator.

Furthermore, this instruction is needed to set up the address for the call. The call itself is only a terminator because it is a tail call. Seems like we would have to have a rather complex set of conditions to decide whether to use a terminator variant or a non-terminator variant.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   ///
   /// @param MO The MachineOperand in question. MO.isReg() should always
   /// be true.
   /// @return Whether this operand is allowed to be used PC-relatively.
   virtual bool isPCRelRegisterOperandLegal(const MachineOperand &MO) const {
     return false;
   }
 
+  /// Allow targets to specify that a particular instruction is part of a
+  /// terminator sequence so that block splitting for the stack smashing
+  /// protector feature does not separate the instruction from the terminator.
+  /// This is for instructions such as the PowerPC MTCTR instruction that
+  /// is clearly not a COPY but sets up the CTR register that a subsequent
+  /// branch uses.
+  virtual bool isPartOfTerminatorSequence(const MachineInstr &MI) const {
+    return false;
+  }
+
 protected:
   /// Target-dependent implementation for IsCopyInstr.
   /// If the specific machine instruction is a instruction that moves/copies
   /// value from one register to another register return destination and source
   /// registers as machine operands.
   virtual Optional<DestSourcePair>
   isCopyInstrImpl(const MachineInstr &MI) const {
     return None;

llvm/lib/CodeGen/CodeGenCommonISel.cpp

do {
       --Previous;
       if (Previous->isCall())
         return SplitPoint;
     } while(Previous->getOpcode() != TII.getCallFrameSetupOpcode());
 
     return Previous;
   }
 
-  while (MIIsInTerminatorSequence(*Previous)) {
+  while (MIIsInTerminatorSequence(*Previous) || TII.isPartOfTerminatorSequence(*Previous)) {
     SplitPoint = Previous;
     if (Previous == Start)
       break;
     --Previous;
   }
 
   return SplitPoint;
 }

llvm/lib/Target/PowerPC/PPCInstrInfo.h

#endif
   MachineInstr *
   findLoopInstr(MachineBasicBlock &PreHeader,
                 SmallPtrSet<MachineBasicBlock *, 8> &Visited) const;
 
   /// Analyze loop L, which must be a single-basic-block loop, and if the
   /// conditions can be understood enough produce a PipelinerLoopInfo object.
   std::unique_ptr<TargetInstrInfo::PipelinerLoopInfo>
   analyzeLoopForPipelining(MachineBasicBlock *LoopBB) const override;
+  bool isPartOfTerminatorSequence(const MachineInstr &MI) const override {
+    return MI.getOpcode() == PPC::MTCTR || MI.getOpcode() == PPC::MTCTR8;
+  }
 };
 
 }
 
 #endif

llvm/test/CodeGen/PowerPC/safestack-indirect-tailcall.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -verify-machineinstrs -mtriple=powerpc64le-unknown-linux-gnu \
+; RUN:   -mcpu=pwr10 -ppc-asm-full-reg-names \
+; RUN:   -ppc-vsr-nums-as-vr < %s | FileCheck %s
+
+%0 = type { %1 }
+%1 = type { { i64, i64 } }
+%2 = type { %3, void (%0*, i32*)* }
+%3 = type { %0, i1 (%0*, %0*, i32)* }
+
+$func2 = comdat any
+
+declare i32 @__gxx_personality_v0(...)
+
+define hidden void @main() local_unnamed_addr #0 personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: main:
+; CHECK:       # %bb.0: # %bb
+; CHECK-NEXT:    mflr r0
+; CHECK-NEXT:    std r0, 16(r1)
+; CHECK-NEXT:    stdu r1, -32(r1)
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    .cfi_offset lr, 16
+; CHECK-NEXT:    paddi r3, 0, func2@PCREL, 1
+; CHECK-NEXT:    std r3, 24(0)
+; CHECK-NEXT:  .Ltmp0:
+; CHECK-NEXT:    bl func1@notoc
+; CHECK-NEXT:  .Ltmp1:
+; CHECK-NEXT:  # %bb.1: # %bb1
+; CHECK-NEXT:  .LBB0_2: # %bb2
+; CHECK-NEXT:  .Ltmp2:
+; CHECK-NEXT:    bl _Unwind_Resume@notoc
+bb:
+  %i = getelementptr inbounds %2, %2* null, i64 0, i32 1
+  store void (%0*, i32*)* @func2, void (%0*, i32*)** %i, align 8
+  invoke void @func1()
+          to label %bb1 unwind label %bb2
+
+bb1:                                              ; preds = %bb
+  unreachable
+
+bb2:                                              ; preds = %bb
+  %i3 = landingpad { i8*, i32 }
+          cleanup
+  resume { i8*, i32 } undef
+}
+
+declare hidden void @func1() local_unnamed_addr #0
+
+; Function Attrs: sspreq
+define hidden void @func2(%0* %arg, i32* %arg1) #1 comdat align 2 {
+; CHECK-LABEL: func2:
+; CHECK:       # %bb.0: # %bb
+; CHECK-NEXT:    mflr r0
+; CHECK-NEXT:    std r0, 16(r1)
+; CHECK-NEXT:    stdu r1, -48(r1)
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    .cfi_offset lr, 16
+; CHECK-NEXT:    ld r4, -28688(r13)
+; CHECK-NEXT:    std r4, 40(r1)
+; CHECK-NEXT:    ld r12, 0(r3)
+; CHECK-NEXT:    ld r3, 40(r1)
+; CHECK-NEXT:    ld r4, -28688(r13)
+; CHECK-NEXT:    cmpld r4, r3
+; CHECK-NEXT:    bne cr0, .LBB1_2
+; CHECK-NEXT:  # %bb.1: # %bb
+; CHECK-NEXT:    mtctr r12
+; CHECK-NEXT:    li r3, 0
+; CHECK-NEXT:    addi r1, r1, 48
+; CHECK-NEXT:    ld r0, 16(r1)
+; CHECK-NEXT:    mtlr r0
+; CHECK-NEXT:    bctr
+; CHECK-NEXT:    #TC_RETURNr8 ctr 0
+; CHECK-NEXT:  .LBB1_2: # %bb
+; CHECK-NEXT:    bl __stack_chk_fail@notoc
+bb:
+  %i = bitcast %0* %arg to void (i32)**
+  %i2 = load void (i32)*, void (i32)** %i, align 8
+  tail call void %i2(i32 noundef signext undef)
+  ret void
+}
+
+attributes #1 = { sspreq }