[llvm] [NVPTX] Fix code generation for `trap-unreachable`. (PR #67478)

[Christian Sigg via llvm-commits]

https://github.com/chsigg created https://github.com/llvm/llvm-project/pull/67478

https://reviews.llvm.org/D152789 added an `exit` op before each `unreachable`. This means we never get to the `trap` instruction.

- When `trap-unreachable` is enabled and `no-trap-after-noreturn` is not, don't insert `exit` before each `unreachable`.
- Lower ISD::TRAP to both `trap` and `exit` instead of just the former.

The fix doesn't work with `no-trap-after-noreturn`, because the `unreachable`s not following a `noreturn` are lowered to `exit; trap; exit;`.

An alternative approach would be to insert `trap`s in `NVPTXLowerUnreachablePass`, depending on the `trap-unreachable` and `no-trap-after-noreturn` settings. I think we would then want skip lowering ISD::TRAP, so that we don't end up with `trap; exit; trap;` sequences.

>From 8506012a3348eb8b4de9222cfc23164bf173ba92 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg at google.com>
Date: Tue, 26 Sep 2023 21:00:33 +0200
Subject: [PATCH] [NVPTX] Fix code generation for `trap-unreachable`.

https://reviews.llvm.org/D152789 added an `exit` op before each `unreachable`. This means we never get to the `trap` instruction.

- When `trap-unreachable` is enabled and `no-trap-after-noreturn` is not, don't insert `exit` before each `unreachable`.
- Lower ISD::TRAP to both `trap` and `exit` instead of just the former.

The fix doesn't work with `no-trap-after-noreturn`, because the `unreachable`s not following a `noreturn` are lowered to `exit; trap; exit;`.

An alternative approach would be to insert `trap`s in `NVPTXLowerUnreachablePass`, depending on the `trap-unreachable` and `no-trap-after-noreturn` settings. I think we would then want skip lowering ISD::TRAP, so that we don't end up with `trap; exit; trap;` sequences.
---
 llvm/lib/Target/NVPTX/NVPTXInstrInfo.td      |  4 +++-
 llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp | 10 ++--------
 llvm/test/CodeGen/NVPTX/unreachable.ll       |  6 +++++-
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td b/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
index ad10d7938ef12e4..fb5ae339cb11a74 100644
--- a/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
+++ b/llvm/lib/Target/NVPTX/NVPTXInstrInfo.td
@@ -3545,7 +3545,9 @@ def Callseq_End :
             [(callseq_end timm:$amt1, timm:$amt2)]>;
 
 // trap instruction
-def trapinst : NVPTXInst<(outs), (ins), "trap;", [(trap)]>;
+// Emit an `exit` as well to convey to ptxas that `trap` exits the CFG.
+// This won't be necessary in a future version of ptxas.
+def trapinst : NVPTXInst<(outs), (ins), "trap; exit;", [(trap)]>;
 
 // Call prototype wrapper
 def SDTCallPrototype : SDTypeProfile<0, 1, [SDTCisInt<0>]>;
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index cad97b1f14eb2b9..0aaeff9fa76b15c 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -63,13 +63,6 @@ static cl::opt<bool> UseShortPointersOpt(
         "Use 32-bit pointers for accessing const/local/shared address spaces."),
     cl::init(false), cl::Hidden);
 
-// FIXME: intended as a temporary debugging aid. Should be removed before it
-// makes it into the LLVM-17 release.
-static cl::opt<bool>
-    ExitOnUnreachable("nvptx-exit-on-unreachable",
-                      cl::desc("Lower 'unreachable' as 'exit' instruction."),
-                      cl::init(true), cl::Hidden);
-
 namespace llvm {
 
 void initializeGenericToNVVMLegacyPassPass(PassRegistry &);
@@ -410,7 +403,8 @@ void NVPTXPassConfig::addIRPasses() {
     addPass(createSROAPass());
   }
 
-  if (ExitOnUnreachable)
+  const auto &options = getNVPTXTargetMachine().Options;
+  if (!options.TrapUnreachable || options.NoTrapAfterNoreturn)
     addPass(createNVPTXLowerUnreachablePass());
 }
 
diff --git a/llvm/test/CodeGen/NVPTX/unreachable.ll b/llvm/test/CodeGen/NVPTX/unreachable.ll
index 742089df1bd4533..05d5a750c444c93 100644
--- a/llvm/test/CodeGen/NVPTX/unreachable.ll
+++ b/llvm/test/CodeGen/NVPTX/unreachable.ll
@@ -1,5 +1,7 @@
 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs -trap-unreachable \
+; RUN:     | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-TRAP
 ; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
 ; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
 
@@ -11,7 +13,9 @@ define void @kernel_func() {
 ; CHECK: call.uni
 ; CHECK: throw,
   call void @throw()
-; CHECK: exit
+; CHECK-TRAP-NOT: exit;
+; CHECK-TRAP: trap;
+; CHECK: exit;
   unreachable
 }