[llvm] [NVPTX] Handle unreachable default in llvm::SwitchIns. (PR #72641)

[via llvm-commits]

https://github.com/mmoadeli created https://github.com/llvm/llvm-project/pull/72641

Addressing [Lower unreachable to exit to allow ptxas to accurately reconstruct the CFG](https://reviews.llvm.org/D152789) resulted in performance regression In scenarios where a Switch Instruction has an unreachable default successor. Substituting the unreachable instruction with an exit instruction introduces an additional block in the Control Flow Graph (CFG), thereby negatively impacting performance. To mitigate this undesirable impact, we refrain from processing blocks that serve as successors to the unreachable default in the switch instruction.These blocks are subsequently optimized out by other passes in the optimization pipeline.

>From 98061eca09243138eed168b0856758ea25a17359 Mon Sep 17 00:00:00 2001
From: m moadeli <mahmoud.moadeli at codeplay.com>
Date: Fri, 17 Nov 2023 12:07:26 +0000
Subject: [PATCH] Handle unreachable default in llvm::SwitchIns.

---
 .../Target/NVPTX/NVPTXLowerUnreachable.cpp    | 26 ++++++++++++++++
 .../CodeGen/NVPTX/unreachable-switch-case.ll  | 31 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 llvm/test/CodeGen/NVPTX/unreachable-switch-case.ll

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
index 34f06b548db260d..0a5e7667f1bbdc0 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
@@ -138,7 +138,32 @@ bool NVPTXLowerUnreachable::runOnFunction(Function &F) {
   InlineAsm *Exit = InlineAsm::get(ExitFTy, "exit;", "", true);
 
   bool Changed = false;
+
+  // In scenarios where a Switch Instruction has an unreachable default
+  // successor, substituting the unreachable instruction with an exit
+  // instruction introduces an additional block in the Control Flow Graph
+  // (CFG), thereby negatively impacting performance. To mitigate this
+  // undesirable impact, we proactively refrain from processing blocks that
+  // serve as successors to the unreachable default in the switch instruction.
+  // It is noteworthy that these blocks are subsequently optimized out by other
+  // passes in the optimization pipeline.
+
+  SmallPtrSet<const BasicBlock *, 4> BlocksToAvoid;
+
   for (auto &BB : F)
+    for (auto &I : BB) {
+      if (auto SI = dyn_cast<SwitchInst>(&I)) {
+        const auto DefaultSuccessorBlock = SI->getDefaultDest();
+        if (DefaultSuccessorBlock->size() == 1 &&
+            dyn_cast<UnreachableInst>(DefaultSuccessorBlock->begin())) {
+          BlocksToAvoid.insert(DefaultSuccessorBlock);
+        }
+      }
+    }
+
+  for (auto &BB : F) {
+    if (BlocksToAvoid.find(&BB) != BlocksToAvoid.end())
+      continue;
     for (auto &I : BB) {
       if (auto unreachableInst = dyn_cast<UnreachableInst>(&I)) {
         if (isLoweredToTrap(*unreachableInst))
@@ -147,6 +172,7 @@ bool NVPTXLowerUnreachable::runOnFunction(Function &F) {
         Changed = true;
       }
     }
+  }
   return Changed;
 }
 
diff --git a/llvm/test/CodeGen/NVPTX/unreachable-switch-case.ll b/llvm/test/CodeGen/NVPTX/unreachable-switch-case.ll
new file mode 100644
index 000000000000000..b2bdffc7eedb605
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/unreachable-switch-case.ll
@@ -0,0 +1,31 @@
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s  
+
+declare noundef i32 @llvm.nvvm.read.ptx.sreg.tid.x() #1
+
+define void @kernel_func() {
+
+  %1 = tail call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+
+  switch i32 %1, label %unreachabledefault [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb1
+    i32 3, label %bb2
+  ]
+
+  bb0:
+    ret void
+
+  bb1:
+    ret void
+
+  bb2:
+    ret void
+
+  unreachabledefault:
+    unreachable
+
+; CHECK:  @kernel_func
+; CHECK-NOT: unreachabledefault
+; CHECK:  -- End function
+}