From f0693277c7cbe982bcfa5d1a8a9540807f6ec314 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Tue, 22 Nov 2022 14:53:43 -0500
Subject: [PATCH] CloneModule: Handling cloning ifuncs

This is tested in a future llvm-reduce patch.
---
 llvm/lib/Transforms/Utils/CloneModule.cpp       | 15 ++++++++++++
 llvm/unittests/Transforms/Utils/CloningTest.cpp | 32 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/llvm/lib/Transforms/Utils/CloneModule.cpp b/llvm/lib/Transforms/Utils/CloneModule.cpp
index 55cda0f..55e0512 100644
--- a/llvm/lib/Transforms/Utils/CloneModule.cpp
+++ b/llvm/lib/Transforms/Utils/CloneModule.cpp
@@ -109,6 +109,15 @@ std::unique_ptr<Module> llvm::CloneModule(
     VMap[&I] = GA;
   }
 
+  for (const GlobalIFunc &I : M.ifuncs()) {
+    // Defer setting the resolver function until after functions are cloned.
+    auto *GI =
+        GlobalIFunc::create(I.getValueType(), I.getAddressSpace(),
+                            I.getLinkage(), I.getName(), nullptr, New.get());
+    GI->copyAttributesFrom(&I);
+    VMap[&I] = GI;
+  }
+
   // Now that all of the things that global variable initializer can refer to
   // have been created, loop through and copy the global variable referrers
   // over...  We also set the attributes on the global now.
@@ -184,6 +193,12 @@ std::unique_ptr<Module> llvm::CloneModule(
       GA->setAliasee(MapValue(C, VMap));
   }
 
+  for (const GlobalIFunc &I : M.ifuncs()) {
+    GlobalIFunc *GI = cast<GlobalIFunc>(VMap[&I]);
+    if (const Constant *Resolver = I.getResolver())
+      GI->setResolver(MapValue(Resolver, VMap));
+  }
+
   // And named metadata....
   for (const NamedMDNode &NMD : M.named_metadata()) {
     NamedMDNode *NewNMD = New->getOrInsertNamedMetadata(NMD.getName());
diff --git a/llvm/unittests/Transforms/Utils/CloningTest.cpp b/llvm/unittests/Transforms/Utils/CloningTest.cpp
index 186714a..9b4a6fb 100644
--- a/llvm/unittests/Transforms/Utils/CloningTest.cpp
+++ b/llvm/unittests/Transforms/Utils/CloningTest.cpp
@@ -922,6 +922,23 @@ protected:
     GV->addMetadata(LLVMContext::MD_type, *MDNode::get(C, {}));
     GV->setComdat(CD);
 
+    // Add ifuncs
+    {
+      const unsigned AddrSpace = 123;
+      auto *FuncPtrTy = Type::getInt8Ty(C)->getPointerTo(123);
+      auto *FuncTy = FunctionType::get(FuncPtrTy, false);
+
+      auto *ResolverF = Function::Create(FuncTy, GlobalValue::PrivateLinkage,
+                                         AddrSpace, "resolver", OldM);
+      BasicBlock *ResolverBody = BasicBlock::Create(C, "", ResolverF);
+      ReturnInst::Create(C, ConstantPointerNull::get(FuncPtrTy), ResolverBody);
+
+      GlobalIFunc *GI = GlobalIFunc::create(FuncTy, AddrSpace,
+                                            GlobalValue::LinkOnceODRLinkage,
+                                            "an_ifunc", ResolverF, OldM);
+      GI->setVisibility(GlobalValue::ProtectedVisibility);
+    }
+
     {
       // Add an empty compile unit first that isn't otherwise referenced, to
       // confirm that compile units get cloned in the correct order.
@@ -1087,4 +1104,19 @@ TEST_F(CloneModule, Comdat) {
   Function *NewF = NewM->getFunction("f");
   EXPECT_EQ(CD, NewF->getComdat());
 }
+
+TEST_F(CloneModule, IFunc) {
+  ASSERT_EQ(1u, NewM->ifunc_size());
+
+  const GlobalIFunc &IFunc = *NewM->ifunc_begin();
+  EXPECT_EQ("an_ifunc", IFunc.getName());
+  EXPECT_EQ(GlobalValue::LinkOnceODRLinkage, IFunc.getLinkage());
+  EXPECT_EQ(GlobalValue::ProtectedVisibility, IFunc.getVisibility());
+  EXPECT_EQ(123u, IFunc.getAddressSpace());
+
+  const Function *Resolver = IFunc.getResolverFunction();
+  ASSERT_NE(nullptr, Resolver);
+  EXPECT_EQ("resolver", Resolver->getName());
+  EXPECT_EQ(GlobalValue::PrivateLinkage, Resolver->getLinkage());
+}
 }
-- 
2.7.4