[llvm] [CodeGen] Avoid generating trap instructions after exception restore intrinsics (PR #109560)

[via llvm-commits]

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

Right now, PrologEpilogInsertion uses `MBB.isReturnBlock()` to determine whether or not to restore CSRs. This normally works with `EH_RETURN` since the latter is marked as being a return instruction; however, the situation gets more complicated when `trap-unreachable` is involved since it's no longer the last instruction in the block. As such, with `trap-unreachable` enabled, PEI misses our block and does not emit a restore when necessary. This also causes a MachineVerifier failure as traps are not considered terminators.

As there are potentially other, similar issues in other passes here, we fix this by marking `llvm.eh.return` as `noreturn` and then adding an additional check in SelectionDAG and GlobalISel before emitting trap instructions when visiting `unreachable`.

>From 3daa4022a395ba939b26fc6f239b89a878456e27 Mon Sep 17 00:00:00 2001
From: duk <37 at cmail.nu>
Date: Sun, 22 Sep 2024 00:29:48 -0400
Subject: [PATCH] [CodeGen] Avoid generating trap instructions after exception
 intrinsics

Right now, PrologEpilogInsertion uses `MBB.isReturnBlock()` to determine whether or not to restore CSRs. This normally works with `EH_RETURN` since the latter is marked as being a return instruction; however, the situation gets more complicated when `trap-unreachable` is involved since it's no longer the last instruction in the block. As such, with `trap-unreachable` enabled, PEI misses our block and does not emit a restore when necessary. This also causes a MachineVerifier failure as traps are not considered terminators.

As there are potentially other, similar issues in other passes here, we fix this by marking `llvm.eh.return` as `noreturn` and then adding an additional check in SelectionDAG and GlobalISel before emitting trap instructions when visiting `unreachable`.
---
 llvm/include/llvm/IR/Intrinsics.td                |  4 ++--
 llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp      |  8 ++++++++
 .../CodeGen/SelectionDAG/SelectionDAGBuilder.cpp  |  8 ++++++++
 llvm/test/CodeGen/X86/eh-trap-unreachable-x86.ll  | 15 +++++++++++++++
 .../CodeGen/X86/eh-trap-unreachable-x86_64.ll     | 15 +++++++++++++++
 5 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/eh-trap-unreachable-x86.ll
 create mode 100644 llvm/test/CodeGen/X86/eh-trap-unreachable-x86_64.ll

diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index 0a74a217a5f010..5e4ee4276aa97f 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -1405,8 +1405,8 @@ let IntrProperties = [IntrNoMem, IntrSpeculatable, IntrWillReturn] in {
 // given function it is 'const' and may be CSE'd etc.
 def int_eh_typeid_for : Intrinsic<[llvm_i32_ty], [llvm_anyptr_ty], [IntrNoMem]>;
 
-def int_eh_return_i32 : Intrinsic<[], [llvm_i32_ty, llvm_ptr_ty]>;
-def int_eh_return_i64 : Intrinsic<[], [llvm_i64_ty, llvm_ptr_ty]>;
+def int_eh_return_i32 : Intrinsic<[], [llvm_i32_ty, llvm_ptr_ty], [IntrNoReturn]>;
+def int_eh_return_i64 : Intrinsic<[], [llvm_i64_ty, llvm_ptr_ty], [IntrNoReturn]>;
 
 // eh.exceptionpointer returns the pointer to the exception caught by
 // the given `catchpad`.
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 8e860a1f740295..7cf6e4d55c24a9 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -3125,6 +3125,14 @@ bool IRTranslator::translateUnreachable(const User &U, MachineIRBuilder &MIRBuil
     // Do not emit an additional trap instruction.
     if (Call->isNonContinuableTrap())
       return true;
+    // Do not emit trap instructions after EH_RETURN intrinsics.
+    // This can cause problems later down the line when other machine passes
+    // attempt to use the last instruction in a BB to determine terminator behavior.
+    if (const auto *II = dyn_cast<IntrinsicInst>(Call)) {
+      const auto IID = II->getIntrinsicID();
+      if (IID == Intrinsic::eh_return_i32 || IID == Intrinsic::eh_return_i64)
+        return true;
+    }
   }
 
   MIRBuilder.buildTrap();
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 25213f587116d5..8ca1aa0278f127 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3553,6 +3553,14 @@ void SelectionDAGBuilder::visitUnreachable(const UnreachableInst &I) {
     // Do not emit an additional trap instruction.
     if (Call->isNonContinuableTrap())
       return;
+    // Do not emit trap instructions after EH_RETURN intrinsics.
+    // This can cause problems later down the line when other machine passes
+    // attempt to use the last instruction in a BB to determine terminator behavior.
+    if (const auto *II = dyn_cast<IntrinsicInst>(Call)) {
+      const auto IID = II->getIntrinsicID();
+      if (IID == Intrinsic::eh_return_i32 || IID == Intrinsic::eh_return_i64)
+        return;
+    }
   }
 
   DAG.setRoot(DAG.getNode(ISD::TRAP, getCurSDLoc(), MVT::Other, DAG.getRoot()));
diff --git a/llvm/test/CodeGen/X86/eh-trap-unreachable-x86.ll b/llvm/test/CodeGen/X86/eh-trap-unreachable-x86.ll
new file mode 100644
index 00000000000000..1101a89f68b0b4
--- /dev/null
+++ b/llvm/test/CodeGen/X86/eh-trap-unreachable-x86.ll
@@ -0,0 +1,15 @@
+; RUN: llc < %s -mtriple=i386-unknown-unknown -trap-unreachable -no-trap-after-noreturn=false -verify-machineinstrs | FileCheck %s
+
+;; Ensure we restore caller-saved registers before EH_RETURN, even with trap-unreachable enabled.
+;; For now, we deliberately avoid generating traps altogether.
+
+; CHECK-LABEL: test32
+; CHECK: pushl
+; CHECK: popl
+; CHECK: eh_return
+; CHECK-NOT: ud2
+define void @test32(i32 %offset, ptr %handler) {
+  call void @llvm.eh.unwind.init()
+  call void @llvm.eh.return.i32(i32 %offset, ptr %handler)
+  unreachable
+}
diff --git a/llvm/test/CodeGen/X86/eh-trap-unreachable-x86_64.ll b/llvm/test/CodeGen/X86/eh-trap-unreachable-x86_64.ll
new file mode 100644
index 00000000000000..d95bd610b5622f
--- /dev/null
+++ b/llvm/test/CodeGen/X86/eh-trap-unreachable-x86_64.ll
@@ -0,0 +1,15 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -trap-unreachable -no-trap-after-noreturn=false -verify-machineinstrs | FileCheck %s
+
+;; Ensure we restore caller-saved registers before EH_RETURN, even with trap-unreachable enabled.
+;; For now, we deliberately avoid generating traps altogether.
+
+; CHECK-LABEL: test64
+; CHECK: pushq
+; CHECK: popq
+; CHECK: eh_return
+; CHECK-NOT: ud2
+define void @test64(i64 %offset, ptr %handler) {
+  call void @llvm.eh.unwind.init()
+  call void @llvm.eh.return.i64(i64 %offset, ptr %handler)
+  unreachable
+}