[llvm] [X86] Avoid invalid ret for EH_RETURN with SHSTK (PR #219322)

[via llvm-commits]

https://github.com/linisha15 updated https://github.com/llvm/llvm-project/pull/219322

>From c2c6691e7420d5b988769a7edcc750902bc4ba16 Mon Sep 17 00:00:00 2001
From: linisha <linisha232 at gmail.com>
Date: Fri, 28 Aug 2026 04:19:59 +0530
Subject: [PATCH 1/3] [X86] Avoid invalid ret for EH_RETURN with SHSTK

---
 llvm/lib/Target/X86/X86ExpandPseudo.cpp  | 12 ++++++++++++
 llvm/test/CodeGen/X86/eh-return-shstk.ll | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/eh-return-shstk.ll

diff --git a/llvm/lib/Target/X86/X86ExpandPseudo.cpp b/llvm/lib/Target/X86/X86ExpandPseudo.cpp
index 612d90361a44e..b8195e9cefc98 100644
--- a/llvm/lib/Target/X86/X86ExpandPseudo.cpp
+++ b/llvm/lib/Target/X86/X86ExpandPseudo.cpp
@@ -402,6 +402,18 @@ bool X86ExpandPseudoImpl::expandMI(MachineBasicBlock &MBB,
     BuildMI(MBB, MBBI, DL,
             TII->get(Uses64BitFramePtr ? X86::MOV64rr : X86::MOV32rr), StackPtr)
         .addReg(DestAddr.getReg());
+    if (STI->hasSHSTK()) {
+      unsigned PopOpcode = Uses64BitFramePtr ? X86::POP64r : X86::POP32r;
+      unsigned JumpOpcode = X86::JMP32r;
+      if (Uses64BitFramePtr)
+        JumpOpcode = STI->isTargetWin64() || STI->isTargetUEFI64()
+                         ? X86::JMP64r_REX
+                         : X86::JMP64r;
+      BuildMI(MBB, MBBI, DL, TII->get(PopOpcode))
+          .addReg(DestAddr.getReg(), RegState::Define);
+      BuildMI(MBB, MBBI, DL, TII->get(JumpOpcode)).addReg(DestAddr.getReg());
+      MBB.erase(MBBI);
+    }
     // The EH_RETURN pseudo is really removed during the MC Lowering.
     return true;
   }
diff --git a/llvm/test/CodeGen/X86/eh-return-shstk.ll b/llvm/test/CodeGen/X86/eh-return-shstk.ll
new file mode 100644
index 0000000000000..5f2c5d194824f
--- /dev/null
+++ b/llvm/test/CodeGen/X86/eh-return-shstk.ll
@@ -0,0 +1,19 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs < %s | FileCheck %s --check-prefix=RET
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+shstk -verify-machineinstrs < %s | FileCheck %s --check-prefix=SHSTK
+
+define void @test(i64 %offset, ptr %handler) {
+entry:
+  call void @llvm.eh.return.i64(i64 %offset, ptr %handler)
+  unreachable
+}
+
+; RET-LABEL: test:
+; RET:       movq %rcx, %rsp
+; RET-NEXT:  retq
+
+; SHSTK-LABEL: test:
+; SHSTK:      movq %rcx, %rsp
+; SHSTK-NEXT: popq %rcx
+; SHSTK-NEXT: jmpq *%rcx
+
+declare void @llvm.eh.return.i64(i64, ptr)

>From 504047b3a47f730de05010f7887dd95dc8816418 Mon Sep 17 00:00:00 2001
From: Linisha <linisha232 at gmail.com>
Date: Mon, 7 Sep 2026 23:07:18 +0530
Subject: [PATCH 2/3] Update llvm/lib/Target/X86/X86ExpandPseudo.cpp

Co-authored-by: Phoebe Wang <phoebe.wang at intel.com>
---
 llvm/lib/Target/X86/X86ExpandPseudo.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86ExpandPseudo.cpp b/llvm/lib/Target/X86/X86ExpandPseudo.cpp
index b8195e9cefc98..bb7b7092fd8bf 100644
--- a/llvm/lib/Target/X86/X86ExpandPseudo.cpp
+++ b/llvm/lib/Target/X86/X86ExpandPseudo.cpp
@@ -403,7 +403,7 @@ bool X86ExpandPseudoImpl::expandMI(MachineBasicBlock &MBB,
             TII->get(Uses64BitFramePtr ? X86::MOV64rr : X86::MOV32rr), StackPtr)
         .addReg(DestAddr.getReg());
     if (STI->hasSHSTK()) {
-      unsigned PopOpcode = Uses64BitFramePtr ? X86::POP64r : X86::POP32r;
+      unsigned PopOpcode = STI->is64Bit() ? X86::POP64r : X86::POP32r;
       unsigned JumpOpcode = X86::JMP32r;
       if (Uses64BitFramePtr)
         JumpOpcode = STI->isTargetWin64() || STI->isTargetUEFI64()

>From 5aaf37a12788a1d059a6f359f0649578a899c6b4 Mon Sep 17 00:00:00 2001
From: linisha <linisha232 at gmail.com>
Date: Mon, 7 Sep 2026 23:33:35 +0530
Subject: [PATCH 3/3] [X86] Add EH_RETURN SHSTK target tests

---
 llvm/test/CodeGen/X86/eh-return-shstk.ll | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/llvm/test/CodeGen/X86/eh-return-shstk.ll b/llvm/test/CodeGen/X86/eh-return-shstk.ll
index 5f2c5d194824f..7281085ee9c72 100644
--- a/llvm/test/CodeGen/X86/eh-return-shstk.ll
+++ b/llvm/test/CodeGen/X86/eh-return-shstk.ll
@@ -1,5 +1,7 @@
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs < %s | FileCheck %s --check-prefix=RET
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+shstk -verify-machineinstrs < %s | FileCheck %s --check-prefix=SHSTK
+; RUN: llc -mtriple=x86_64-pc-windows-msvc -mattr=+shstk -verify-machineinstrs < %s | FileCheck %s --check-prefix=WIN64
+; RUN: llc -mtriple=x86_64-uefi -mattr=+shstk -verify-machineinstrs < %s | FileCheck %s --check-prefix=UEFI64
 
 define void @test(i64 %offset, ptr %handler) {
 entry:
@@ -16,4 +18,14 @@ entry:
 ; SHSTK-NEXT: popq %rcx
 ; SHSTK-NEXT: jmpq *%rcx
 
+; WIN64-LABEL: test:
+; WIN64:       movq %rcx, %rsp
+; WIN64-NEXT:  popq %rcx
+; WIN64-NEXT:  rex64 jmpq *%rcx
+
+; UEFI64-LABEL: test:
+; UEFI64:       movq %rcx, %rsp
+; UEFI64-NEXT:  popq %rcx
+; UEFI64-NEXT:  rex64 jmpq *%rcx
+
 declare void @llvm.eh.return.i64(i64, ptr)