[CodeGen] Generalize trap emission after SP check fail (#109744)

Generalize and improve some target-specific code that emits traps after
stack protector failure in SelectionDAG & GlobalIsel.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -3913,16 +3913,11 @@ bool IRTranslator::emitSPDescriptorFailure(StackProtectorDescriptor &SPD,
     return false;
   }
 
-  // On PS4/PS5, the "return address" must still be within the calling
-  // function, even if it's at the very end, so emit an explicit TRAP here.
-  // WebAssembly needs an unreachable instruction after a non-returning call,
-  // because the function return type can be different from __stack_chk_fail's
-  // return type (void).
-  const TargetMachine &TM = MF->getTarget();
-  if (TM.getTargetTriple().isPS() || TM.getTargetTriple().isWasm()) {
-    LLVM_DEBUG(dbgs() << "Unhandled trap emission for stack protector fail\n");
-    return false;
-  }
+  // Emit a trap instruction if we are required to do so.
+  const TargetOptions &TargetOpts = TLI->getTargetMachine().Options;
+  if (TargetOpts.TrapUnreachable && !TargetOpts.NoTrapAfterNoreturn)
+    CurBuilder->buildInstr(TargetOpcode::G_TRAP);
+
   return true;
 }
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -3183,15 +3183,10 @@ SelectionDAGBuilder::visitSPDescriptorFailure(StackProtectorDescriptor &SPD) {
   SDValue Chain = TLI.makeLibCall(DAG, RTLIB::STACKPROTECTOR_CHECK_FAIL,
                                   MVT::isVoid, {}, CallOptions, getCurSDLoc())
                       .second;
-  // On PS4/PS5, the "return address" must still be within the calling
-  // function, even if it's at the very end, so emit an explicit TRAP here.
-  // Passing 'true' for doesNotReturn above won't generate the trap for us.
-  if (TM.getTargetTriple().isPS())
-    Chain = DAG.getNode(ISD::TRAP, getCurSDLoc(), MVT::Other, Chain);
-  // WebAssembly needs an unreachable instruction after a non-returning call,
-  // because the function return type can be different from __stack_chk_fail's
-  // return type (void).
-  if (TM.getTargetTriple().isWasm())
+
+  // Emit a trap instruction if we are required to do so.
+  const TargetOptions &TargetOpts = DAG.getTarget().Options;
+  if (TargetOpts.TrapUnreachable && !TargetOpts.NoTrapAfterNoreturn)
     Chain = DAG.getNode(ISD::TRAP, getCurSDLoc(), MVT::Other, Chain);
 
   DAG.setRoot(Chain);

llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

@@ -133,7 +133,10 @@ WebAssemblyTargetMachine::WebAssemblyTargetMachine(
   // WebAssembly type-checks instructions, but a noreturn function with a return
   // type that doesn't match the context will cause a check failure. So we lower
   // LLVM 'unreachable' to ISD::TRAP and then lower that to WebAssembly's
-  // 'unreachable' instructions which is meant for that case.
+  // 'unreachable' instructions which is meant for that case. Formerly, we also
+  // needed to add checks to SP failure emission in the instruction selection
+  // backends, but this has since been tied to TrapUnreachable and is no longer
+  // necessary.
   this->Options.TrapUnreachable = true;
   this->Options.NoTrapAfterNoreturn = false;
 

llvm/lib/Target/X86/X86TargetMachine.cpp

@@ -242,7 +242,10 @@ X86TargetMachine::X86TargetMachine(const Target &T, const Triple &TT,
           OL),
       TLOF(createTLOF(getTargetTriple())), IsJIT(JIT) {
   // On PS4/PS5, the "return address" of a 'noreturn' call must still be within
-  // the calling function, and TrapUnreachable is an easy way to get that.
+  // the calling function. Note that this also includes __stack_chk_fail,
+  // so there was some target-specific logic in the instruction selectors
+  // to handle that. That code has since been generalized, so the only thing
+  // needed is to set TrapUnreachable here.
   if (TT.isPS() || TT.isOSBinFormatMachO()) {
     this->Options.TrapUnreachable = true;
     this->Options.NoTrapAfterNoreturn = TT.isOSBinFormatMachO();

llvm/test/CodeGen/AArch64/GlobalISel/irtranslator-stack-protector-trap-unreachable.ll

@@ -0,0 +1,80 @@
+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+; RUN: llc -mtriple=aarch64 -global-isel -stop-after=irtranslator -verify-machineinstrs \
+; RUN:     -trap-unreachable=false -o - %s | FileCheck -check-prefix=NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -global-isel -stop-after=irtranslator -verify-machineinstrs \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o - %s | FileCheck -check-prefix=NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -global-isel -stop-after=irtranslator -verify-machineinstrs \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o - %s | FileCheck -check-prefix=TRAP_UNREACHABLE %s
+
+;; Make sure we emit trap instructions after stack protector checks iff NoTrapAfterNoReturn is false.
+
+define void @test() nounwind ssp {
+  ; NO_TRAP_UNREACHABLE-LABEL: name: test
+  ; NO_TRAP_UNREACHABLE: bb.1.entry:
+  ; NO_TRAP_UNREACHABLE-NEXT:   successors: %bb.2(0x7ffff800), %bb.3(0x00000800)
+  ; NO_TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0.StackGuardSlot
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD:%[0-9]+]]:gpr64sp(p0) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD1:%[0-9]+]]:gpr64sp(p0) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; NO_TRAP_UNREACHABLE-NEXT:   G_STORE [[LOAD_STACK_GUARD1]](p0), [[FRAME_INDEX]](p0) :: (volatile store (p0) into %stack.0.StackGuardSlot)
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX1:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.1.buf
+  ; NO_TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
+  ; NO_TRAP_UNREACHABLE-NEXT:   $x0 = COPY [[FRAME_INDEX1]](p0)
+  ; NO_TRAP_UNREACHABLE-NEXT:   BL @callee, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit-def $w0
+  ; NO_TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX2:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0.StackGuardSlot
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[LOAD:%[0-9]+]]:_(s64) = G_LOAD [[FRAME_INDEX2]](p0) :: (volatile load (s64) from %stack.0.StackGuardSlot)
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD2:%[0-9]+]]:gpr64sp(s64) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; NO_TRAP_UNREACHABLE-NEXT:   [[ICMP:%[0-9]+]]:_(s1) = G_ICMP intpred(ne), [[LOAD_STACK_GUARD2]](s64), [[LOAD]]
+  ; NO_TRAP_UNREACHABLE-NEXT:   G_BRCOND [[ICMP]](s1), %bb.3
+  ; NO_TRAP_UNREACHABLE-NEXT:   G_BR %bb.2
+  ; NO_TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; NO_TRAP_UNREACHABLE-NEXT: bb.3.entry:
+  ; NO_TRAP_UNREACHABLE-NEXT:   successors:
+  ; NO_TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; NO_TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
+  ; NO_TRAP_UNREACHABLE-NEXT:   BL &__stack_chk_fail, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
+  ; NO_TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
+  ; NO_TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; NO_TRAP_UNREACHABLE-NEXT: bb.2.entry:
+  ; NO_TRAP_UNREACHABLE-NEXT:   RET_ReallyLR
+  ;
+  ; TRAP_UNREACHABLE-LABEL: name: test
+  ; TRAP_UNREACHABLE: bb.1.entry:
+  ; TRAP_UNREACHABLE-NEXT:   successors: %bb.2(0x7ffff800), %bb.3(0x00000800)
+  ; TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0.StackGuardSlot
+  ; TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD:%[0-9]+]]:gpr64sp(p0) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD1:%[0-9]+]]:gpr64sp(p0) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; TRAP_UNREACHABLE-NEXT:   G_STORE [[LOAD_STACK_GUARD1]](p0), [[FRAME_INDEX]](p0) :: (volatile store (p0) into %stack.0.StackGuardSlot)
+  ; TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX1:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.1.buf
+  ; TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
+  ; TRAP_UNREACHABLE-NEXT:   $x0 = COPY [[FRAME_INDEX1]](p0)
+  ; TRAP_UNREACHABLE-NEXT:   BL @callee, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit-def $w0
+  ; TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
+  ; TRAP_UNREACHABLE-NEXT:   [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+  ; TRAP_UNREACHABLE-NEXT:   [[FRAME_INDEX2:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0.StackGuardSlot
+  ; TRAP_UNREACHABLE-NEXT:   [[LOAD:%[0-9]+]]:_(s64) = G_LOAD [[FRAME_INDEX2]](p0) :: (volatile load (s64) from %stack.0.StackGuardSlot)
+  ; TRAP_UNREACHABLE-NEXT:   [[LOAD_STACK_GUARD2:%[0-9]+]]:gpr64sp(s64) = LOAD_STACK_GUARD :: (dereferenceable invariant load (p0) from @__stack_chk_guard)
+  ; TRAP_UNREACHABLE-NEXT:   [[ICMP:%[0-9]+]]:_(s1) = G_ICMP intpred(ne), [[LOAD_STACK_GUARD2]](s64), [[LOAD]]
+  ; TRAP_UNREACHABLE-NEXT:   G_BRCOND [[ICMP]](s1), %bb.3
+  ; TRAP_UNREACHABLE-NEXT:   G_BR %bb.2
+  ; TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; TRAP_UNREACHABLE-NEXT: bb.3.entry:
+  ; TRAP_UNREACHABLE-NEXT:   successors:
+  ; TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
+  ; TRAP_UNREACHABLE-NEXT:   BL &__stack_chk_fail, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
+  ; TRAP_UNREACHABLE-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
+  ; TRAP_UNREACHABLE-NEXT:   G_TRAP
+  ; TRAP_UNREACHABLE-NEXT: {{  $}}
+  ; TRAP_UNREACHABLE-NEXT: bb.2.entry:
+  ; TRAP_UNREACHABLE-NEXT:   RET_ReallyLR
+entry:
+  %buf = alloca [8 x i8]
+  %result = call i32(ptr) @callee(ptr %buf) nounwind
+  ret void
+}
+
+declare i32 @callee(ptr) nounwind

llvm/test/CodeGen/AArch64/stack-protector-trap-unreachable.ll

@@ -0,0 +1,33 @@
+;; Make sure we emit trap instructions after stack protector checks iff NoTrapAfterNoReturn is false.
+
+; RUN: llc -mtriple=aarch64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,TRAP_UNREACHABLE %s
+
+;; Make sure FastISel doesn't break anything.
+; RUN: llc -mtriple=aarch64 -fast-isel -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -fast-isel -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=aarch64 -fast-isel -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,TRAP_UNREACHABLE %s
+
+; CHECK-LABEL: Machine code for function test
+; CHECK: bb.0.entry:
+; CHECK:  BL {{.}}__stack_chk_fail
+; CHECK-NEXT: ADJCALLSTACKUP
+; TRAP_UNREACHABLE-NEXT: BRK 1
+; NO_TRAP_UNREACHABLE-NOT: BRK 1
+; NO_TRAP_UNREACHABLE-EMPTY:
+
+define void @test() nounwind ssp {
+entry:
+  %buf = alloca [8 x i8]
+  %result = call i32(ptr) @callee(ptr %buf) nounwind
+  ret void
+}
+
+declare i32 @callee(ptr) nounwind

llvm/test/CodeGen/X86/stack-protector-trap-unreachable.ll

@@ -0,0 +1,40 @@
+;; Make sure we emit trap instructions after stack protector checks iff NoTrapAfterNoReturn is false.
+
+; RUN: llc -enable-selectiondag-sp -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -enable-selectiondag-sp -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -enable-selectiondag-sp -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,TRAP_UNREACHABLE %s
+
+; RUN: llc -enable-selectiondag-sp=false -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -enable-selectiondag-sp=false -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -enable-selectiondag-sp=false -mtriple=x86_64 -fast-isel=false -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,TRAP_UNREACHABLE %s
+
+;; Make sure FastISel doesn't break anything.
+; RUN: llc -mtriple=x86_64 -fast-isel -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable=false --o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=x86_64 -fast-isel -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,NO_TRAP_UNREACHABLE %s
+; RUN: llc -mtriple=x86_64 -fast-isel -global-isel=false -verify-machineinstrs -print-after=finalize-isel \
+; RUN:     -trap-unreachable -no-trap-after-noreturn=false -o /dev/null 2>&1 %s | FileCheck --check-prefixes=CHECK,TRAP_UNREACHABLE %s
+
+; CHECK-LABEL: Machine code for function test
+; CHECK: bb.0.entry:
+; CHECK: CALL64{{.*}}__stack_chk_fail
+; CHECK-NEXT: ADJCALLSTACKUP64
+; TRAP_UNREACHABLE-NEXT: TRAP
+; NO_TRAP_UNREACHABLE-NOT: TRAP
+; NO_TRAP_UNREACHABLE-EMPTY:
+
+define void @test() nounwind ssp {
+entry:
+  %buf = alloca [8 x i8]
+  %result = call i32(ptr) @callee(ptr %buf) nounwind
+  ret void
+}
+
+declare i32 @callee(ptr) nounwind

llvm/test/CodeGen/X86/unreachable-trap.ll

@@ -1,7 +1,15 @@
 ; RUN: llc -o - %s -mtriple=x86_64-linux-gnu | FileCheck %s --check-prefixes=CHECK
 ; RUN: llc -o - %s -mtriple=x86_64-windows-msvc | FileCheck %s --check-prefixes=CHECK
-; RUN: llc -o - %s -mtriple=x86_64-scei-ps4 | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
 ; RUN: llc -o - %s -mtriple=x86_64-apple-darwin | FileCheck %s --check-prefixes=CHECK,NO_TRAP_AFTER_NORETURN
+
+; On PS4/PS5, always emit trap instructions regardless of of trap-unreachable or no-trap-after-noreturn.
+; RUN: llc -o - %s -mtriple=x86_64-scei-ps4 -trap-unreachable | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+; RUN: llc -o - %s -mtriple=x86_64-sie-ps5 -trap-unreachable | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+; RUN: llc -o - %s -mtriple=x86_64-scei-ps4 -trap-unreachable=false | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+; RUN: llc -o - %s -mtriple=x86_64-sie-ps5 -trap-unreachable=false | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+; RUN: llc -o - %s -mtriple=x86_64-scei-ps4 -trap-unreachable -no-trap-after-noreturn=false | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+; RUN: llc -o - %s -mtriple=x86_64-sie-ps5 -trap-unreachable -no-trap-after-noreturn=false | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
+
 ; RUN: llc --trap-unreachable -o - %s -mtriple=x86_64-linux-gnu | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
 ; RUN: llc --trap-unreachable -global-isel -o - %s -mtriple=x86_64-linux-gnu | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN
 ; RUN: llc --trap-unreachable -fast-isel -o - %s -mtriple=x86_64-linux-gnu | FileCheck %s --check-prefixes=CHECK,TRAP_AFTER_NORETURN