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

[via llvm-commits]

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

>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 1/2] 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 34f06b548db260..0a5e7667f1bbdc 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 00000000000000..b2bdffc7eedb60
--- /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
+}

>From 70fa78e0102bdcc6cccbb787b982b7ae6738eb7b Mon Sep 17 00:00:00 2001
From: M Moadeli <mahmoud.moadeli at codeplay.com>
Date: Thu, 15 Feb 2024 09:54:13 +0000
Subject: [PATCH 2/2] Skip processing BasicBlocks with single unreachable
 instruction.

Updates the `nvptx-lower-unreachable` pass to bypass BasicBlocks containing just one unreachable instruction. This allows
further optimization by subsequent passes, enhancing overall performance and efficiency.
---
 .../Target/NVPTX/NVPTXLowerUnreachable.cpp    | 31 ++++++-------------
 1 file changed, 9 insertions(+), 22 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
index 0a5e7667f1bbdc..be70fa51b2f6b7 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerUnreachable.cpp
@@ -139,30 +139,17 @@ bool NVPTXLowerUnreachable::runOnFunction(Function &F) {
 
   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);
-        }
-      }
-    }
+  // In scenarios where a BasicBlock contains only one unreachable instruction,
+  // the joint action of nvptx-isel and unreachable-mbb-elimination
+  // effectively optimizes the BasicBlock out. However, adding an exit
+  // command to such a BasicBlock, as suggested by this pass, preserves it
+  // within the Control Flow Graph (CFG), thereby negatively impacting size and
+  // performance. To counteract this undesirable consequence, we choose to
+  // refrain from processing BasicBlocks with just one unreachable instruction
+  // in this pass.
 
   for (auto &BB : F) {
-    if (BlocksToAvoid.find(&BB) != BlocksToAvoid.end())
+    if ((BB.size() == 1) && (isa<UnreachableInst>(BB.front())))
       continue;
     for (auto &I : BB) {
       if (auto unreachableInst = dyn_cast<UnreachableInst>(&I)) {