Index: lib/Transforms/Utils/CodeExtractor.cpp =================================================================== --- lib/Transforms/Utils/CodeExtractor.cpp +++ lib/Transforms/Utils/CodeExtractor.cpp @@ -871,6 +871,15 @@ if (BFI && NumExitBlocks > 1) calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI); + // If some blocks are captured by a block address, remap to new Function + for (auto &BB : *newFunction) { + SmallVector Users(BB.user_begin(), BB.user_end()); + for (auto *U : Users) { + if (auto *BA = dyn_cast(U)) + BA->replaceAllUsesWith(BlockAddress::get(newFunction, &BB)); + } + } + // Loop over all of the PHI nodes in the header block, and change any // references to the old incoming edge to be the new incoming edge. for (BasicBlock::iterator I = header->begin(); isa(I); ++I) { Index: test/Transforms/CodeExtractor/BlockAddressreference.ll =================================================================== --- test/Transforms/CodeExtractor/BlockAddressreference.ll +++ test/Transforms/CodeExtractor/BlockAddressreference.ll @@ -0,0 +1,33 @@ +; RUN: opt < %s -loop-extract -disable-output + +@label = common local_unnamed_addr global i8* null + +define i32 @sterix(i32 %n) { +entry: + %tobool = icmp ne i32 %n, 0 + ; this blockaddress references a basic block that goes in the extracted loop + %cond = select i1 %tobool, i8* blockaddress(@sterix, %for.cond), i8* blockaddress(@sterix, %exit) + store i8* %cond, i8** @label + %cmp5 = icmp sgt i32 %n, 0 + br i1 %cmp5, label %for.body, label %exit + +for.cond: + %mul = shl nsw i32 %s.06, 1 + %exitcond = icmp eq i32 %inc, %n + br i1 %exitcond, label %exit.loopexit, label %for.body + +for.body: + %i.07 = phi i32 [ %inc, %for.cond ], [ 0, %entry ] + %s.06 = phi i32 [ %mul, %for.cond ], [ 1, %entry ] + %inc = add nuw nsw i32 %i.07, 1 + br label %for.cond + +exit.loopexit: + %phitmp = icmp ne i32 %s.06, 2 + %phitmp8 = zext i1 %phitmp to i32 + br label %exit + +exit: + %s.1 = phi i32 [ 1, %entry ], [ %phitmp8, %exit.loopexit ] + ret i32 %s.1 +}