Index: lib/Transforms/Utils/SplitModule.cpp =================================================================== --- lib/Transforms/Utils/SplitModule.cpp +++ lib/Transforms/Utils/SplitModule.cpp @@ -20,6 +20,7 @@ #include "llvm/ADT/Hashing.h" #include "llvm/ADT/MapVector.h" #include "llvm/ADT/SetVector.h" +#include "llvm/IR/Constants.h" #include "llvm/IR/Function.h" #include "llvm/IR/GlobalAlias.h" #include "llvm/IR/GlobalObject.h" @@ -55,6 +56,24 @@ } } +// Adds all GlobalValue users of V to the same cluster as GV. +static void addAllGlobalValueUsers(ClusterMapType &GVtoClusterMap, + const GlobalValue *GV, const Value *V) { + for (auto *U : V->users()) { + SmallVector Worklist; + Worklist.push_back(U); + while (!Worklist.empty()) { + const User *UU = Worklist.pop_back_val(); + // For each constant that is not a GV (a pure const) recurse. + if (isa(UU) && !isa(UU)) { + Worklist.append(UU->user_begin(), UU->user_end()); + continue; + } + addNonConstUser(GVtoClusterMap, GV, UU); + } + } +} + // Find partitions for module in the way that no locals need to be // globalized. // Try to balance pack those partitions into N files since this roughly equals @@ -64,7 +83,6 @@ // At this point module should have the proper mix of globals and locals. // As we attempt to partition this module, we must not change any // locals to globals. - DEBUG(dbgs() << "Partition module with (" << M->size() << ")functions\n"); ClusterMapType GVtoClusterMap; ComdatMembersType ComdatMembers; @@ -95,23 +113,17 @@ GVtoClusterMap.unionSets(&GV, Base); } - // Further only iterate over local GVs. - if (!GV.hasLocalLinkage()) - return; - - for (auto *U : GV.users()) { - SmallVector Worklist; - Worklist.push_back(U); - while (!Worklist.empty()) { - const User *UU = Worklist.pop_back_val(); - // For each constant that is not a GV (a pure const) recurse. - if (isa(UU) && !isa(UU)) { - Worklist.append(UU->user_begin(), UU->user_end()); + if (const Function *F = dyn_cast(&GV)) { + for (const BasicBlock &BB : *F) { + BlockAddress *BA = BlockAddress::lookup(&BB); + if (!BA || !BA->isConstantUsed()) continue; - } - addNonConstUser(GVtoClusterMap, &GV, UU); + addAllGlobalValueUsers(GVtoClusterMap, F, BA); } } + + if (GV.hasLocalLinkage()) + addAllGlobalValueUsers(GVtoClusterMap, &GV, &GV); }; std::for_each(M->begin(), M->end(), recordGVSet); @@ -225,37 +237,26 @@ externalize(&GV); for (GlobalAlias &GA : M->aliases()) externalize(&GA); + } - // FIXME: We should be able to reuse M as the last partition instead of - // cloning it. - for (unsigned I = 0; I != N; ++I) { - ValueToValueMapTy VMap; - std::unique_ptr MPart( - CloneModule(M.get(), VMap, [=](const GlobalValue *GV) { - return isInPartition(GV, I, N); - })); - if (I != 0) - MPart->setModuleInlineAsm(""); - ModuleCallback(std::move(MPart)); - } - } else { - // This performs splitting without a need for externalization, which might not - // always be possible. - ClusterIDMapType ClusterIDMap; - findPartitions(M.get(), ClusterIDMap, N); + // This performs splitting without a need for externalization, which might not + // always be possible. + ClusterIDMapType ClusterIDMap; + findPartitions(M.get(), ClusterIDMap, N); - for (unsigned I = 0; I < N; ++I) { - ValueToValueMapTy VMap; - std::unique_ptr MPart( - CloneModule(M.get(), VMap, [&](const GlobalValue *GV) { - if (ClusterIDMap.count(GV)) - return (ClusterIDMap[GV] == I); - else - return isInPartition(GV, I, N); - })); - if (I != 0) - MPart->setModuleInlineAsm(""); - ModuleCallback(std::move(MPart)); - } + // FIXME: We should be able to reuse M as the last partition instead of + // cloning it. + for (unsigned I = 0; I < N; ++I) { + ValueToValueMapTy VMap; + std::unique_ptr MPart( + CloneModule(M.get(), VMap, [&](const GlobalValue *GV) { + if (ClusterIDMap.count(GV)) + return (ClusterIDMap[GV] == I); + else + return isInPartition(GV, I, N); + })); + if (I != 0) + MPart->setModuleInlineAsm(""); + ModuleCallback(std::move(MPart)); } } Index: test/tools/llvm-split/blockaddress.ll =================================================================== --- /dev/null +++ test/tools/llvm-split/blockaddress.ll @@ -0,0 +1,37 @@ +; Test that blockaddress target is in the same partition. +; RUN: llvm-split -j5 -o %t %s +; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s +; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1234 %s +; RUN: llvm-dis -o - %t2 | FileCheck --check-prefix=CHECK1234 %s +; RUN: llvm-dis -o - %t3 | FileCheck --check-prefix=CHECK1234 %s +; RUN: llvm-dis -o - %t4 | FileCheck --check-prefix=CHECK1234 %s + +; CHECK0: @xxx = global [2 x i8*] [i8* blockaddress(@f, %exit), i8* blockaddress(@g, %exit)] +; CHECK1234: @xxx = external global [2 x i8*] +; CHECK1234-NOT: blockaddress +@xxx = global [2 x i8*] [i8* blockaddress(@f, %exit), i8* blockaddress(@g, %exit)] + +; CHECK0: define i32 @f() +; CHECK1234: declare i32 @f() +define i32 @f(){ +entry: + br label %exit +exit: + ret i32 0 +} + +; CHECK0: define i32 @g() +; CHECK1234: declare i32 @g() +define i32 @g(){ +entry: + br label %exit +exit: + ret i32 0 +} + +; CHECK0: define i8* @h() +; CHECK1234: declare i8* @h() +define i8* @h(){ +entry: + ret i8* blockaddress(@f, %exit) +}