[llvm] [NVPTX] Fix NVPTXLowerUnreachable::isLoweredToTrap logic (PR #109730)

[via llvm-commits]

https://github.com/duk-37 created https://github.com/llvm/llvm-project/pull/109730

Previously, this pass would not generate traps if `NoTrapAfterNoreturn` was set and would generate traps even if the instruction directly before the `UnreachableInst` was `llvm.trap()`.

Fix both of these problems and add some tests.

>From eee13ba115d3016f46b3e21bd7a1eb6037b43b99 Mon Sep 17 00:00:00 2001
From: duk <37 at cmail.nu>
Date: Mon, 23 Sep 2024 19:11:37 -0400
Subject: [PATCH] [NVPTX] Fix NVPTXLowerUnreachable::isLoweredToTrap logic

Previously, this pass would not generate traps if `NoTrapAfterNoreturn` was set and would generate traps even if the instruction directly before the `UnreachableInst` was `llvm.trap()`.
---
 .../Target/NVPTX/NVPTXLowerUnreachable.cpp    | 19 +++++++++-----
 llvm/test/CodeGen/NVPTX/unreachable.ll        | 26 +++++++++++++++----
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
index 92b90e2559154e..ef0d0e7cd7d510 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
@@ -110,17 +110,24 @@ StringRef NVPTXLowerUnreachable::getPassName() const {
 }
 
 // =============================================================================
-// Returns whether a `trap` intrinsic should be emitted before I.
+// Returns whether a `trap` intrinsic would be emitted before I.
 //
 // This is a copy of the logic in SelectionDAGBuilder::visitUnreachable().
 // =============================================================================
 bool NVPTXLowerUnreachable::isLoweredToTrap(const UnreachableInst &I) const {
-  if (!TrapUnreachable)
-    return false;
-  if (!NoTrapAfterNoreturn)
-    return true;
   const CallInst *Call = dyn_cast_or_null<CallInst>(I.getPrevNode());
-  return Call && Call->doesNotReturn();
+  if (Call) {
+    // We've already emitted a non-continuable trap.
+    if (Call->isNonContinuableTrap())
+      return true;
+
+    // No traps are emitted for calls that do not return when this option is enabled.
+    if (NoTrapAfterNoreturn && Call->doesNotReturn())
+      return false;
+  }
+
+  // In all other cases, we will generate a trap if TrapUnreachable is set.
+  return TrapUnreachable;
 }
 
 // =============================================================================
diff --git a/llvm/test/CodeGen/NVPTX/unreachable.ll b/llvm/test/CodeGen/NVPTX/unreachable.ll
index 011497c4e23401..f9118900cb7372 100644
--- a/llvm/test/CodeGen/NVPTX/unreachable.ll
+++ b/llvm/test/CodeGen/NVPTX/unreachable.ll
@@ -1,18 +1,23 @@
-; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs \
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs -trap-unreachable=false \
 ; RUN:     | FileCheck %s  --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
-; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs \
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs -trap-unreachable=false \
 ; RUN:     | FileCheck %s  --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
-; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs -trap-unreachable \
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs -trap-unreachable -no-trap-after-noreturn \
+; RUN:     | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs -trap-unreachable -no-trap-after-noreturn \
+; RUN:     | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs -trap-unreachable -no-trap-after-noreturn=false \
 ; RUN:     | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-TRAP
-; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs -trap-unreachable \
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs -trap-unreachable -no-trap-after-noreturn=false \
 ; 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 %}
 
 ; CHECK: .extern .func throw
 declare void @throw() #0
+declare void @llvm.trap() #0
 
-; CHECK: .entry kernel_func
+; CHECK-LABEL: .entry kernel_func
 define void @kernel_func() {
 ; CHECK: call.uni
 ; CHECK: throw,
@@ -24,6 +29,17 @@ define void @kernel_func() {
   unreachable
 }
 
+; CHECK-LABEL: kernel_func_2
+define void @kernel_func_2() {
+; CHECK: trap; exit;
+  call void @llvm.trap()
+
+;; Make sure we avoid emitting two trap instructions.
+; CHECK-NOT: trap;
+; CHECK-NOT: exit;
+  unreachable
+}
+
 attributes #0 = { noreturn }