Implement the logic in NVPTXLowerUnreachablePass instead.

Author: chsigg

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -3226,14 +3226,9 @@ void SelectionDAGBuilder::visitUnreachable(const UnreachableInst &I) {
 
   // We may be able to ignore unreachable behind a noreturn call.
   if (DAG.getTarget().Options.NoTrapAfterNoreturn) {
-    const BasicBlock &BB = *I.getParent();
-    if (&I != &BB.front()) {
-      BasicBlock::const_iterator PredI =
-        std::prev(BasicBlock::const_iterator(&I));
-      if (const CallInst *Call = dyn_cast<CallInst>(&*PredI)) {
+    if (const CallInst *Call = dyn_cast_or_null<CallInst>(I.getPrevNode())) {
         if (Call->doesNotReturn())
           return;
-      }
     }
   }
 

llvm/lib/Target/NVPTX/NVPTX.h

@@ -47,7 +47,8 @@ MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
 FunctionPass *createNVPTXLowerArgsPass();
 FunctionPass *createNVPTXLowerAllocaPass();
-FunctionPass *createNVPTXLowerUnreachablePass();
+FunctionPass *createNVPTXLowerUnreachablePass(bool TrapUnreachable,
+                                              bool NoTrapAfterNoreturn);
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 

llvm/lib/Target/NVPTX/NVPTXInstrInfo.td

@@ -3545,9 +3545,7 @@ def Callseq_End :
             [(callseq_end timm:$amt1, timm:$amt2)]>;
 
 // trap instruction
-// 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)]>;
+def trapinst : NVPTXInst<(outs), (ins), "trap;", [(trap)]>;
 
 // Call prototype wrapper
 def SDTCallPrototype : SDTypeProfile<0, 1, [SDTCisInt<0>]>;

llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp

@@ -72,6 +72,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Pass.h"
 
@@ -83,14 +84,19 @@ void initializeNVPTXLowerUnreachablePass(PassRegistry &);
 
 namespace {
 class NVPTXLowerUnreachable : public FunctionPass {
+  StringRef getPassName() const override;
   bool runOnFunction(Function &F) override;
+  bool shouldEmitTrap(const UnreachableInst &I) const;
 
 public:
   static char ID; // Pass identification, replacement for typeid
-  NVPTXLowerUnreachable() : FunctionPass(ID) {}
-  StringRef getPassName() const override {
-    return "add an exit instruction before every unreachable";
-  }
+  NVPTXLowerUnreachable(bool TrapUnreachable, bool NoTrapAfterNoreturn)
+      : FunctionPass(ID), TrapUnreachable(TrapUnreachable),
+        NoTrapAfterNoreturn(NoTrapAfterNoreturn) {}
+
+private:
+  bool TrapUnreachable;
+  bool NoTrapAfterNoreturn;
 };
 } // namespace
 
@@ -99,6 +105,24 @@ char NVPTXLowerUnreachable::ID = 1;
 INITIALIZE_PASS(NVPTXLowerUnreachable, "nvptx-lower-unreachable",
                 "Lower Unreachable", false, false)
 
+StringRef NVPTXLowerUnreachable::getPassName() const {
+  return "add an exit instruction before every unreachable";
+}
+
+// =============================================================================
+// Returns whether a `trap` intrinsic should be emitted before I.
+//
+// This is a copy of the logic in SelectionDAGBuilder::visitUnreachable().
+// =============================================================================
+bool NVPTXLowerUnreachable::shouldEmitTrap(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();
+}
+
 // =============================================================================
 // Main function for this pass.
 // =============================================================================
@@ -109,18 +133,25 @@ bool NVPTXLowerUnreachable::runOnFunction(Function &F) {
   LLVMContext &C = F.getContext();
   FunctionType *ExitFTy = FunctionType::get(Type::getVoidTy(C), false);
   InlineAsm *Exit = InlineAsm::get(ExitFTy, "exit;", "", true);
+  Function *Trap = nullptr;
 
   bool Changed = false;
   for (auto &BB : F)
     for (auto &I : BB) {
       if (auto unreachableInst = dyn_cast<UnreachableInst>(&I)) {
-        Changed = true;
+        if (shouldEmitTrap(*unreachableInst)) {
+          if (!Trap)
+            Trap = Intrinsic::getDeclaration(F.getParent(), Intrinsic::trap);
+          CallInst::Create(Trap, "", unreachableInst);
+        }
         CallInst::Create(ExitFTy, Exit, "", unreachableInst);
+        Changed = true;
       }
     }
   return Changed;
 }
 
-FunctionPass *llvm::createNVPTXLowerUnreachablePass() {
-  return new NVPTXLowerUnreachable();
+FunctionPass *llvm::createNVPTXLowerUnreachablePass(bool TrapUnreachable,
+                                                    bool NoTrapAfterNoreturn) {
+  return new NVPTXLowerUnreachable(TrapUnreachable, NoTrapAfterNoreturn);
 }

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -403,9 +403,9 @@ void NVPTXPassConfig::addIRPasses() {
     addPass(createSROAPass());
   }
 
-  const auto &options = getNVPTXTargetMachine().Options;
-  if (!options.TrapUnreachable || options.NoTrapAfterNoreturn)
-    addPass(createNVPTXLowerUnreachablePass());
+  const auto &Options = getNVPTXTargetMachine().Options;
+  addPass(createNVPTXLowerUnreachablePass(Options.TrapUnreachable,
+                                          Options.NoTrapAfterNoreturn));
 }
 
 bool NVPTXPassConfig::addInstSelector() {

llvm/test/CodeGen/NVPTX/unreachable.ll

@@ -1,5 +1,9 @@
-; 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=nvptx -mcpu=sm_20 -verify-machineinstrs \
+; RUN:     | FileCheck %s  --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs \
+; RUN:     | FileCheck %s  --check-prefix=CHECK --check-prefix=CHECK-NOTRAP
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs -trap-unreachable \
+; RUN:     | FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-TRAP
 ; 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 %}
@@ -15,6 +19,7 @@ define void @kernel_func() {
   call void @throw()
 ; CHECK-TRAP-NOT: exit;
 ; CHECK-TRAP: trap;
+; CHECK-NOTRAP-NOT: trap;
 ; CHECK: exit;
   unreachable
 }