1 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "llvm/Transforms/Coroutines/CoroElide.h" 10 #include "CoroInternal.h" 11 #include "llvm/ADT/DenseMap.h" 12 #include "llvm/ADT/Statistic.h" 13 #include "llvm/Analysis/AliasAnalysis.h" 14 #include "llvm/Analysis/InstructionSimplify.h" 15 #include "llvm/Analysis/OptimizationRemarkEmitter.h" 16 #include "llvm/IR/Dominators.h" 17 #include "llvm/IR/InstIterator.h" 18 #include "llvm/Support/ErrorHandling.h" 19 #include "llvm/Support/FileSystem.h" 20 #include <optional> 21 22 using namespace llvm; 23 24 #define DEBUG_TYPE "coro-elide" 25 26 STATISTIC(NumOfCoroElided, "The # of coroutine get elided."); 27 28 #ifndef NDEBUG 29 static cl::opt<std::string> CoroElideInfoOutputFilename( 30 "coro-elide-info-output-file", cl::value_desc("filename"), 31 cl::desc("File to record the coroutines got elided"), cl::Hidden); 32 #endif 33 34 namespace { 35 // Created on demand if the coro-elide pass has work to do. 36 struct Lowerer : coro::LowererBase { 37 SmallVector<CoroIdInst *, 4> CoroIds; 38 SmallVector<CoroBeginInst *, 1> CoroBegins; 39 SmallVector<CoroAllocInst *, 1> CoroAllocs; 40 SmallVector<CoroSubFnInst *, 4> ResumeAddr; 41 DenseMap<CoroBeginInst *, SmallVector<CoroSubFnInst *, 4>> DestroyAddr; 42 SmallPtrSet<const SwitchInst *, 4> CoroSuspendSwitches; 43 44 Lowerer(Module &M) : LowererBase(M) {} 45 46 void elideHeapAllocations(Function *F, uint64_t FrameSize, Align FrameAlign, 47 AAResults &AA); 48 bool shouldElide(Function *F, DominatorTree &DT) const; 49 void collectPostSplitCoroIds(Function *F); 50 bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT, 51 OptimizationRemarkEmitter &ORE); 52 bool hasEscapePath(const CoroBeginInst *, 53 const SmallPtrSetImpl<BasicBlock *> &) const; 54 }; 55 } // end anonymous namespace 56 57 // Go through the list of coro.subfn.addr intrinsics and replace them with the 58 // provided constant. 59 static void replaceWithConstant(Constant *Value, 60 SmallVectorImpl<CoroSubFnInst *> &Users) { 61 if (Users.empty()) 62 return; 63 64 // See if we need to bitcast the constant to match the type of the intrinsic 65 // being replaced. Note: All coro.subfn.addr intrinsics return the same type, 66 // so we only need to examine the type of the first one in the list. 67 Type *IntrTy = Users.front()->getType(); 68 Type *ValueTy = Value->getType(); 69 if (ValueTy != IntrTy) { 70 // May need to tweak the function type to match the type expected at the 71 // use site. 72 assert(ValueTy->isPointerTy() && IntrTy->isPointerTy()); 73 Value = ConstantExpr::getBitCast(Value, IntrTy); 74 } 75 76 // Now the value type matches the type of the intrinsic. Replace them all! 77 for (CoroSubFnInst *I : Users) 78 replaceAndRecursivelySimplify(I, Value); 79 } 80 81 // See if any operand of the call instruction references the coroutine frame. 82 static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) { 83 for (Value *Op : CI->operand_values()) 84 if (!AA.isNoAlias(Op, Frame)) 85 return true; 86 return false; 87 } 88 89 // Look for any tail calls referencing the coroutine frame and remove tail 90 // attribute from them, since now coroutine frame resides on the stack and tail 91 // call implies that the function does not references anything on the stack. 92 // However if it's a musttail call, we cannot remove the tailcall attribute. 93 // It's safe to keep it there as the musttail call is for symmetric transfer, 94 // and by that point the frame should have been destroyed and hence not 95 // interfering with operands. 96 static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) { 97 Function &F = *Frame->getFunction(); 98 for (Instruction &I : instructions(F)) 99 if (auto *Call = dyn_cast<CallInst>(&I)) 100 if (Call->isTailCall() && operandReferences(Call, Frame, AA) && 101 !Call->isMustTailCall()) 102 Call->setTailCall(false); 103 } 104 105 // Given a resume function @f.resume(%f.frame* %frame), returns the size 106 // and expected alignment of %f.frame type. 107 static std::optional<std::pair<uint64_t, Align>> 108 getFrameLayout(Function *Resume) { 109 // Pull information from the function attributes. 110 auto Size = Resume->getParamDereferenceableBytes(0); 111 if (!Size) 112 return std::nullopt; 113 return std::make_pair(Size, Resume->getParamAlign(0).valueOrOne()); 114 } 115 116 // Finds first non alloca instruction in the entry block of a function. 117 static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) { 118 for (Instruction &I : F->getEntryBlock()) 119 if (!isa<AllocaInst>(&I)) 120 return &I; 121 llvm_unreachable("no terminator in the entry block"); 122 } 123 124 #ifndef NDEBUG 125 static std::unique_ptr<raw_fd_ostream> getOrCreateLogFile() { 126 assert(!CoroElideInfoOutputFilename.empty() && 127 "coro-elide-info-output-file shouldn't be empty"); 128 std::error_code EC; 129 auto Result = std::make_unique<raw_fd_ostream>(CoroElideInfoOutputFilename, 130 EC, sys::fs::OF_Append); 131 if (!EC) 132 return Result; 133 llvm::errs() << "Error opening coro-elide-info-output-file '" 134 << CoroElideInfoOutputFilename << " for appending!\n"; 135 return std::make_unique<raw_fd_ostream>(2, false); // stderr. 136 } 137 #endif 138 139 // To elide heap allocations we need to suppress code blocks guarded by 140 // llvm.coro.alloc and llvm.coro.free instructions. 141 void Lowerer::elideHeapAllocations(Function *F, uint64_t FrameSize, 142 Align FrameAlign, AAResults &AA) { 143 LLVMContext &C = F->getContext(); 144 auto *InsertPt = 145 getFirstNonAllocaInTheEntryBlock(CoroIds.front()->getFunction()); 146 147 // Replacing llvm.coro.alloc with false will suppress dynamic 148 // allocation as it is expected for the frontend to generate the code that 149 // looks like: 150 // id = coro.id(...) 151 // mem = coro.alloc(id) ? malloc(coro.size()) : 0; 152 // coro.begin(id, mem) 153 auto *False = ConstantInt::getFalse(C); 154 for (auto *CA : CoroAllocs) { 155 CA->replaceAllUsesWith(False); 156 CA->eraseFromParent(); 157 } 158 159 // FIXME: Design how to transmit alignment information for every alloca that 160 // is spilled into the coroutine frame and recreate the alignment information 161 // here. Possibly we will need to do a mini SROA here and break the coroutine 162 // frame into individual AllocaInst recreating the original alignment. 163 const DataLayout &DL = F->getParent()->getDataLayout(); 164 auto FrameTy = ArrayType::get(Type::getInt8Ty(C), FrameSize); 165 auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt); 166 Frame->setAlignment(FrameAlign); 167 auto *FrameVoidPtr = 168 new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt); 169 170 for (auto *CB : CoroBegins) { 171 CB->replaceAllUsesWith(FrameVoidPtr); 172 CB->eraseFromParent(); 173 } 174 175 // Since now coroutine frame lives on the stack we need to make sure that 176 // any tail call referencing it, must be made non-tail call. 177 removeTailCallAttribute(Frame, AA); 178 } 179 180 bool Lowerer::hasEscapePath(const CoroBeginInst *CB, 181 const SmallPtrSetImpl<BasicBlock *> &TIs) const { 182 const auto &It = DestroyAddr.find(CB); 183 assert(It != DestroyAddr.end()); 184 185 // Limit the number of blocks we visit. 186 unsigned Limit = 32 * (1 + It->second.size()); 187 188 SmallVector<const BasicBlock *, 32> Worklist; 189 Worklist.push_back(CB->getParent()); 190 191 SmallPtrSet<const BasicBlock *, 32> Visited; 192 // Consider basicblock of coro.destroy as visited one, so that we 193 // skip the path pass through coro.destroy. 194 for (auto *DA : It->second) 195 Visited.insert(DA->getParent()); 196 197 do { 198 const auto *BB = Worklist.pop_back_val(); 199 if (!Visited.insert(BB).second) 200 continue; 201 if (TIs.count(BB)) 202 return true; 203 204 // Conservatively say that there is potentially a path. 205 if (!--Limit) 206 return true; 207 208 auto TI = BB->getTerminator(); 209 // Although the default dest of coro.suspend switches is suspend pointer 210 // which means a escape path to normal terminator, it is reasonable to skip 211 // it since coroutine frame doesn't change outside the coroutine body. 212 if (isa<SwitchInst>(TI) && 213 CoroSuspendSwitches.count(cast<SwitchInst>(TI))) { 214 Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(1)); 215 Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(2)); 216 } else 217 Worklist.append(succ_begin(BB), succ_end(BB)); 218 219 } while (!Worklist.empty()); 220 221 // We have exhausted all possible paths and are certain that coro.begin can 222 // not reach to any of terminators. 223 return false; 224 } 225 226 bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const { 227 // If no CoroAllocs, we cannot suppress allocation, so elision is not 228 // possible. 229 if (CoroAllocs.empty()) 230 return false; 231 232 // Check that for every coro.begin there is at least one coro.destroy directly 233 // referencing the SSA value of that coro.begin along each 234 // non-exceptional path. 235 // If the value escaped, then coro.destroy would have been referencing a 236 // memory location storing that value and not the virtual register. 237 238 SmallPtrSet<BasicBlock *, 8> Terminators; 239 // First gather all of the non-exceptional terminators for the function. 240 // Consider the final coro.suspend as the real terminator when the current 241 // function is a coroutine. 242 for (BasicBlock &B : *F) { 243 auto *TI = B.getTerminator(); 244 if (TI->getNumSuccessors() == 0 && !TI->isExceptionalTerminator() && 245 !isa<UnreachableInst>(TI)) 246 Terminators.insert(&B); 247 } 248 249 // Filter out the coro.destroy that lie along exceptional paths. 250 SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins; 251 for (const auto &It : DestroyAddr) { 252 // If there is any coro.destroy dominates all of the terminators for the 253 // coro.begin, we could know the corresponding coro.begin wouldn't escape. 254 for (Instruction *DA : It.second) { 255 if (llvm::all_of(Terminators, [&](auto *TI) { 256 return DT.dominates(DA, TI->getTerminator()); 257 })) { 258 ReferencedCoroBegins.insert(It.first); 259 break; 260 } 261 } 262 263 // Whether there is any paths from coro.begin to Terminators which not pass 264 // through any of the coro.destroys. 265 // 266 // hasEscapePath is relatively slow, so we avoid to run it as much as 267 // possible. 268 if (!ReferencedCoroBegins.count(It.first) && 269 !hasEscapePath(It.first, Terminators)) 270 ReferencedCoroBegins.insert(It.first); 271 } 272 273 // If size of the set is the same as total number of coro.begin, that means we 274 // found a coro.free or coro.destroy referencing each coro.begin, so we can 275 // perform heap elision. 276 return ReferencedCoroBegins.size() == CoroBegins.size(); 277 } 278 279 void Lowerer::collectPostSplitCoroIds(Function *F) { 280 CoroIds.clear(); 281 CoroSuspendSwitches.clear(); 282 for (auto &I : instructions(F)) { 283 if (auto *CII = dyn_cast<CoroIdInst>(&I)) 284 if (CII->getInfo().isPostSplit()) 285 // If it is the coroutine itself, don't touch it. 286 if (CII->getCoroutine() != CII->getFunction()) 287 CoroIds.push_back(CII); 288 289 // Consider case like: 290 // %0 = call i8 @llvm.coro.suspend(...) 291 // switch i8 %0, label %suspend [i8 0, label %resume 292 // i8 1, label %cleanup] 293 // and collect the SwitchInsts which are used by escape analysis later. 294 if (auto *CSI = dyn_cast<CoroSuspendInst>(&I)) 295 if (CSI->hasOneUse() && isa<SwitchInst>(CSI->use_begin()->getUser())) { 296 SwitchInst *SWI = cast<SwitchInst>(CSI->use_begin()->getUser()); 297 if (SWI->getNumCases() == 2) 298 CoroSuspendSwitches.insert(SWI); 299 } 300 } 301 } 302 303 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA, 304 DominatorTree &DT, OptimizationRemarkEmitter &ORE) { 305 CoroBegins.clear(); 306 CoroAllocs.clear(); 307 ResumeAddr.clear(); 308 DestroyAddr.clear(); 309 310 // Collect all coro.begin and coro.allocs associated with this coro.id. 311 for (User *U : CoroId->users()) { 312 if (auto *CB = dyn_cast<CoroBeginInst>(U)) 313 CoroBegins.push_back(CB); 314 else if (auto *CA = dyn_cast<CoroAllocInst>(U)) 315 CoroAllocs.push_back(CA); 316 } 317 318 // Collect all coro.subfn.addrs associated with coro.begin. 319 // Note, we only devirtualize the calls if their coro.subfn.addr refers to 320 // coro.begin directly. If we run into cases where this check is too 321 // conservative, we can consider relaxing the check. 322 for (CoroBeginInst *CB : CoroBegins) { 323 for (User *U : CB->users()) 324 if (auto *II = dyn_cast<CoroSubFnInst>(U)) 325 switch (II->getIndex()) { 326 case CoroSubFnInst::ResumeIndex: 327 ResumeAddr.push_back(II); 328 break; 329 case CoroSubFnInst::DestroyIndex: 330 DestroyAddr[CB].push_back(II); 331 break; 332 default: 333 llvm_unreachable("unexpected coro.subfn.addr constant"); 334 } 335 } 336 337 // PostSplit coro.id refers to an array of subfunctions in its Info 338 // argument. 339 ConstantArray *Resumers = CoroId->getInfo().Resumers; 340 assert(Resumers && "PostSplit coro.id Info argument must refer to an array" 341 "of coroutine subfunctions"); 342 auto *ResumeAddrConstant = 343 Resumers->getAggregateElement(CoroSubFnInst::ResumeIndex); 344 345 replaceWithConstant(ResumeAddrConstant, ResumeAddr); 346 347 bool ShouldElide = shouldElide(CoroId->getFunction(), DT); 348 if (!ShouldElide) 349 ORE.emit([&]() { 350 if (auto FrameSizeAndAlign = 351 getFrameLayout(cast<Function>(ResumeAddrConstant))) 352 return OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId) 353 << "'" << ore::NV("callee", CoroId->getCoroutine()->getName()) 354 << "' not elided in '" 355 << ore::NV("caller", CoroId->getFunction()->getName()) 356 << "' (frame_size=" 357 << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align=" 358 << ore::NV("align", FrameSizeAndAlign->second.value()) << ")"; 359 else 360 return OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId) 361 << "'" << ore::NV("callee", CoroId->getCoroutine()->getName()) 362 << "' not elided in '" 363 << ore::NV("caller", CoroId->getFunction()->getName()) 364 << "' (frame_size=unknown, align=unknown)"; 365 }); 366 367 auto *DestroyAddrConstant = Resumers->getAggregateElement( 368 ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex); 369 370 for (auto &It : DestroyAddr) 371 replaceWithConstant(DestroyAddrConstant, It.second); 372 373 if (ShouldElide) { 374 if (auto FrameSizeAndAlign = 375 getFrameLayout(cast<Function>(ResumeAddrConstant))) { 376 elideHeapAllocations(CoroId->getFunction(), FrameSizeAndAlign->first, 377 FrameSizeAndAlign->second, AA); 378 coro::replaceCoroFree(CoroId, /*Elide=*/true); 379 NumOfCoroElided++; 380 #ifndef NDEBUG 381 if (!CoroElideInfoOutputFilename.empty()) 382 *getOrCreateLogFile() 383 << "Elide " << CoroId->getCoroutine()->getName() << " in " 384 << CoroId->getFunction()->getName() << "\n"; 385 #endif 386 ORE.emit([&]() { 387 return OptimizationRemark(DEBUG_TYPE, "CoroElide", CoroId) 388 << "'" << ore::NV("callee", CoroId->getCoroutine()->getName()) 389 << "' elided in '" 390 << ore::NV("caller", CoroId->getFunction()->getName()) 391 << "' (frame_size=" 392 << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align=" 393 << ore::NV("align", FrameSizeAndAlign->second.value()) << ")"; 394 }); 395 } else { 396 ORE.emit([&]() { 397 return OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId) 398 << "'" << ore::NV("callee", CoroId->getCoroutine()->getName()) 399 << "' not elided in '" 400 << ore::NV("caller", CoroId->getFunction()->getName()) 401 << "' (frame_size=unknown, align=unknown)"; 402 }); 403 } 404 } 405 406 return true; 407 } 408 409 static bool declaresCoroElideIntrinsics(Module &M) { 410 return coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.id.async"}); 411 } 412 413 PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) { 414 auto &M = *F.getParent(); 415 if (!declaresCoroElideIntrinsics(M)) 416 return PreservedAnalyses::all(); 417 418 Lowerer L(M); 419 L.CoroIds.clear(); 420 L.collectPostSplitCoroIds(&F); 421 // If we did not find any coro.id, there is nothing to do. 422 if (L.CoroIds.empty()) 423 return PreservedAnalyses::all(); 424 425 AAResults &AA = AM.getResult<AAManager>(F); 426 DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F); 427 auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F); 428 429 bool Changed = false; 430 for (auto *CII : L.CoroIds) 431 Changed |= L.processCoroId(CII, AA, DT, ORE); 432 433 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); 434 } 435