From 0d527e56a5ee3fd0435e6f99a10692e5115f36b7 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Tue, 22 Nov 2022 21:02:20 -0500
Subject: [PATCH] GlobalIFunc: Make ifunc respect function address spaces

---
 llvm/lib/Bitcode/Reader/BitcodeReader.cpp      |  4 ++--
 llvm/lib/IR/Verifier.cpp                       |  2 +-
 llvm/test/Assembler/ifunc-program-addrspace.ll | 33 ++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Assembler/ifunc-program-addrspace.ll

diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index c480560..198fe33 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -2873,8 +2873,8 @@ Error BitcodeReader::resolveGlobalAndIndirectSymbolInits() {
         Type *ResolverFTy =
             GlobalIFunc::getResolverFunctionType(GI->getValueType());
         // Transparently fix up the type for compatibility with older bitcode
-        GI->setResolver(
-            ConstantExpr::getBitCast(C, ResolverFTy->getPointerTo()));
+        GI->setResolver(ConstantExpr::getBitCast(
+            C, ResolverFTy->getPointerTo(GI->getAddressSpace())));
       } else {
         return error("Expected an alias or an ifunc");
       }
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 002b521..ce0d4da 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -889,7 +889,7 @@ void Verifier::visitGlobalIFunc(const GlobalIFunc &GI) {
   const Type *ResolverTy = GI.getResolver()->getType();
   const Type *ResolverFuncTy =
       GlobalIFunc::getResolverFunctionType(GI.getValueType());
-  Check(ResolverTy == ResolverFuncTy->getPointerTo(),
+  Check(ResolverTy == ResolverFuncTy->getPointerTo(GI.getAddressSpace()),
         "IFunc resolver has incorrect type", &GI);
 }
 
diff --git a/llvm/test/Assembler/ifunc-program-addrspace.ll b/llvm/test/Assembler/ifunc-program-addrspace.ll
new file mode 100644
index 0000000..8933812
--- /dev/null
+++ b/llvm/test/Assembler/ifunc-program-addrspace.ll
@@ -0,0 +1,33 @@
+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
+
+target datalayout = "P1"
+
+; CHECK: @ifunc_as0 = ifunc void (), ptr @resolver_as0
+@ifunc_as0 = ifunc void (), ptr @resolver_as0
+
+; CHECK: @ifunc_as1 = ifunc void (), ptr addrspace(1) @resolver_as1
+@ifunc_as1 = ifunc void (), ptr addrspace(1) @resolver_as1
+
+; CHECK: define ptr @resolver_as0() addrspace(0) {
+define ptr @resolver_as0() addrspace(0) {
+  ret ptr null
+}
+
+; CHECK: define ptr @resolver_as1() addrspace(1) {
+define ptr @resolver_as1() addrspace(1) {
+  ret ptr null
+}
+
+; CHECK: define void @call_ifunc_as0() addrspace(1) {
+; CHECK-NEXT: call addrspace(0) void @ifunc_as0()
+define void @call_ifunc_as0() addrspace(1) {
+  call addrspace(0) void @ifunc_as0()
+  ret void
+}
+
+; CHECK: define void @call_ifunc_as1() addrspace(1) {
+; CHECK-NEXT: call addrspace(1) void @ifunc_as1()
+define void @call_ifunc_as1() addrspace(1) {
+  call addrspace(1) void @ifunc_as1()
+  ret void
+}
-- 
2.7.4