[llvm] [SPIR-V] Patch invoke/callbr call sites when legalizing function signatures (PR #217000)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/217000

>From 58adb9d01ad2c8ef772dc77e4ca0239c54e21918 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 18 Aug 2026 13:44:50 +0200
Subject: [PATCH] [SPIR-V] Patch invoke/callbr call sites when legalizing
 function signatures

Only CallInst was patched before RAUW, leaving invoke/callbr with a stale FunctionType and a broken verifier
---
 .../Target/SPIRV/SPIRVPrepareFunctions.cpp    |  8 +++++--
 .../passes/SPIRVPrepareFunctions-invoke.ll    | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/passes/SPIRVPrepareFunctions-invoke.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp b/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
index d784de3de3610..e8355031f6587 100644
--- a/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
@@ -637,9 +637,13 @@ SPIRVPrepareFunctionsImpl::removeAggregateTypesFromSignature(Function *F) {
       std::move(ChangedTypes), NewF->getName());
 
   for (User *U : F->users()) {
-    if (auto *CI = dyn_cast<CallInst>(U); CI && CI->getCalledFunction() == F)
-      CI->mutateFunctionType(NewF->getFunctionType());
+    if (auto *CB = dyn_cast<CallBase>(U); CB && CB->getCalledFunction() == F)
+      CB->mutateFunctionType(NewF->getFunctionType());
   }
+  // NewF keeps F's address space, so their pointer types match and
+  // RAUW is safe despite the differing signatures.
+  assert(F->getType() == NewF->getType() &&
+         "RAUW requires F and NewF to share the same pointer type");
   F->replaceAllUsesWith(NewF);
 
   // register the mutation
diff --git a/llvm/test/CodeGen/SPIRV/passes/SPIRVPrepareFunctions-invoke.ll b/llvm/test/CodeGen/SPIRV/passes/SPIRVPrepareFunctions-invoke.ll
new file mode 100644
index 0000000000000..c2abe7501f79b
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/passes/SPIRVPrepareFunctions-invoke.ll
@@ -0,0 +1,23 @@
+; RUN: opt -S -passes=spirv-prepare-functions -mtriple=spirv64-unknown-unknown < %s | FileCheck %s
+
+declare i32 @__gxx_personality_v0(...)
+
+; invoke call sites must get their FunctionType patched too, like plain calls.
+; CHECK-LABEL: define void @invoke_caller(
+; CHECK: invoke i32 @callback(i32 %x)
+define void @invoke_caller({ float, float } %x) personality ptr @__gxx_personality_v0 {
+entry:
+  %r = invoke { float, float } @callback({ float, float } %x)
+          to label %cont unwind label %lpad
+cont:
+  ret void
+lpad:
+  %lp = landingpad { ptr, i32 }
+          cleanup
+  ret void
+}
+
+; CHECK-LABEL: define i32 @callback(
+define { float, float } @callback({ float, float } %x) {
+  ret { float, float } %x
+}