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
22using namespace llvm;
23
24#define DEBUG_TYPE "coro-elide"
25
26STATISTIC(NumOfCoroElided, "The # of coroutine get elided.");
27
28#ifndef NDEBUG
29static 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
34namespace {
35// Created on demand if the coro-elide pass has work to do.
36struct 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.
59static 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.
82static 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.
96static 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.
107static std::optional<std::pair<uint64_t, Align>>
108getFrameLayout(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.
117static 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
125static 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.
141void 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, PointerType::getUnqual(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
180bool 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  SmallPtrSet<const BasicBlock *, 32> EscapingBBs;
198  for (auto *U : CB->users()) {
199    // The use from coroutine intrinsics are not a problem.
200    if (isa<CoroFreeInst, CoroSubFnInst, CoroSaveInst>(U))
201      continue;
202
203    // Think all other usages may be an escaping candidate conservatively.
204    //
205    // Note that the major user of switch ABI coroutine (the C++) will store
206    // resume.fn, destroy.fn and the index to the coroutine frame immediately.
207    // So the parent of the coro.begin in C++ will be always escaping.
208    // Then we can't get any performance benefits for C++ by improving the
209    // precision of the method.
210    //
211    // The reason why we still judge it is we want to make LLVM Coroutine in
212    // switch ABIs to be self contained as much as possible instead of a
213    // by-product of C++20 Coroutines.
214    EscapingBBs.insert(cast<Instruction>(U)->getParent());
215  }
216
217  bool PotentiallyEscaped = false;
218
219  do {
220    const auto *BB = Worklist.pop_back_val();
221    if (!Visited.insert(BB).second)
222      continue;
223
224    // A Path insensitive marker to test whether the coro.begin escapes.
225    // It is intentional to make it path insensitive while it may not be
226    // precise since we don't want the process to be too slow.
227    PotentiallyEscaped |= EscapingBBs.count(BB);
228
229    if (TIs.count(BB)) {
230      if (isa<ReturnInst>(BB->getTerminator()) || PotentiallyEscaped)
231        return true;
232
233      // If the function ends with the exceptional terminator, the memory used
234      // by the coroutine frame can be released by stack unwinding
235      // automatically. So we can think the coro.begin doesn't escape if it
236      // exits the function by exceptional terminator.
237
238      continue;
239    }
240
241    // Conservatively say that there is potentially a path.
242    if (!--Limit)
243      return true;
244
245    auto TI = BB->getTerminator();
246    // Although the default dest of coro.suspend switches is suspend pointer
247    // which means a escape path to normal terminator, it is reasonable to skip
248    // it since coroutine frame doesn't change outside the coroutine body.
249    if (isa<SwitchInst>(TI) &&
250        CoroSuspendSwitches.count(cast<SwitchInst>(TI))) {
251      Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(1));
252      Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(2));
253    } else
254      Worklist.append(succ_begin(BB), succ_end(BB));
255
256  } while (!Worklist.empty());
257
258  // We have exhausted all possible paths and are certain that coro.begin can
259  // not reach to any of terminators.
260  return false;
261}
262
263bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const {
264  // If no CoroAllocs, we cannot suppress allocation, so elision is not
265  // possible.
266  if (CoroAllocs.empty())
267    return false;
268
269  // Check that for every coro.begin there is at least one coro.destroy directly
270  // referencing the SSA value of that coro.begin along each
271  // non-exceptional path.
272  // If the value escaped, then coro.destroy would have been referencing a
273  // memory location storing that value and not the virtual register.
274
275  SmallPtrSet<BasicBlock *, 8> Terminators;
276  // First gather all of the terminators for the function.
277  // Consider the final coro.suspend as the real terminator when the current
278  // function is a coroutine.
279  for (BasicBlock &B : *F) {
280    auto *TI = B.getTerminator();
281
282    if (TI->getNumSuccessors() != 0 || isa<UnreachableInst>(TI))
283      continue;
284
285    Terminators.insert(&B);
286  }
287
288  // Filter out the coro.destroy that lie along exceptional paths.
289  SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
290  for (const auto &It : DestroyAddr) {
291    // If every terminators is dominated by coro.destroy, we could know the
292    // corresponding coro.begin wouldn't escape.
293    //
294    // Otherwise hasEscapePath would decide whether there is any paths from
295    // coro.begin to Terminators which not pass through any of the
296    // coro.destroys.
297    //
298    // hasEscapePath is relatively slow, so we avoid to run it as much as
299    // possible.
300    if (llvm::all_of(Terminators,
301                     [&](auto *TI) {
302                       return llvm::any_of(It.second, [&](auto *DA) {
303                         return DT.dominates(DA, TI->getTerminator());
304                       });
305                     }) ||
306        !hasEscapePath(It.first, Terminators))
307      ReferencedCoroBegins.insert(It.first);
308  }
309
310  // If size of the set is the same as total number of coro.begin, that means we
311  // found a coro.free or coro.destroy referencing each coro.begin, so we can
312  // perform heap elision.
313  return ReferencedCoroBegins.size() == CoroBegins.size();
314}
315
316void Lowerer::collectPostSplitCoroIds(Function *F) {
317  CoroIds.clear();
318  CoroSuspendSwitches.clear();
319  for (auto &I : instructions(F)) {
320    if (auto *CII = dyn_cast<CoroIdInst>(&I))
321      if (CII->getInfo().isPostSplit())
322        // If it is the coroutine itself, don't touch it.
323        if (CII->getCoroutine() != CII->getFunction())
324          CoroIds.push_back(CII);
325
326    // Consider case like:
327    // %0 = call i8 @llvm.coro.suspend(...)
328    // switch i8 %0, label %suspend [i8 0, label %resume
329    //                              i8 1, label %cleanup]
330    // and collect the SwitchInsts which are used by escape analysis later.
331    if (auto *CSI = dyn_cast<CoroSuspendInst>(&I))
332      if (CSI->hasOneUse() && isa<SwitchInst>(CSI->use_begin()->getUser())) {
333        SwitchInst *SWI = cast<SwitchInst>(CSI->use_begin()->getUser());
334        if (SWI->getNumCases() == 2)
335          CoroSuspendSwitches.insert(SWI);
336      }
337  }
338}
339
340bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
341                            DominatorTree &DT, OptimizationRemarkEmitter &ORE) {
342  CoroBegins.clear();
343  CoroAllocs.clear();
344  ResumeAddr.clear();
345  DestroyAddr.clear();
346
347  // Collect all coro.begin and coro.allocs associated with this coro.id.
348  for (User *U : CoroId->users()) {
349    if (auto *CB = dyn_cast<CoroBeginInst>(U))
350      CoroBegins.push_back(CB);
351    else if (auto *CA = dyn_cast<CoroAllocInst>(U))
352      CoroAllocs.push_back(CA);
353  }
354
355  // Collect all coro.subfn.addrs associated with coro.begin.
356  // Note, we only devirtualize the calls if their coro.subfn.addr refers to
357  // coro.begin directly. If we run into cases where this check is too
358  // conservative, we can consider relaxing the check.
359  for (CoroBeginInst *CB : CoroBegins) {
360    for (User *U : CB->users())
361      if (auto *II = dyn_cast<CoroSubFnInst>(U))
362        switch (II->getIndex()) {
363        case CoroSubFnInst::ResumeIndex:
364          ResumeAddr.push_back(II);
365          break;
366        case CoroSubFnInst::DestroyIndex:
367          DestroyAddr[CB].push_back(II);
368          break;
369        default:
370          llvm_unreachable("unexpected coro.subfn.addr constant");
371        }
372  }
373
374  // PostSplit coro.id refers to an array of subfunctions in its Info
375  // argument.
376  ConstantArray *Resumers = CoroId->getInfo().Resumers;
377  assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
378                     "of coroutine subfunctions");
379  auto *ResumeAddrConstant =
380      Resumers->getAggregateElement(CoroSubFnInst::ResumeIndex);
381
382  replaceWithConstant(ResumeAddrConstant, ResumeAddr);
383
384  bool ShouldElide = shouldElide(CoroId->getFunction(), DT);
385  if (!ShouldElide)
386    ORE.emit([&]() {
387      if (auto FrameSizeAndAlign =
388              getFrameLayout(cast<Function>(ResumeAddrConstant)))
389        return OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId)
390               << "'" << ore::NV("callee", CoroId->getCoroutine()->getName())
391               << "' not elided in '"
392               << ore::NV("caller", CoroId->getFunction()->getName())
393               << "' (frame_size="
394               << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align="
395               << ore::NV("align", FrameSizeAndAlign->second.value()) << ")";
396      else
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  auto *DestroyAddrConstant = Resumers->getAggregateElement(
405      ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
406
407  for (auto &It : DestroyAddr)
408    replaceWithConstant(DestroyAddrConstant, It.second);
409
410  if (ShouldElide) {
411    if (auto FrameSizeAndAlign =
412            getFrameLayout(cast<Function>(ResumeAddrConstant))) {
413      elideHeapAllocations(CoroId->getFunction(), FrameSizeAndAlign->first,
414                           FrameSizeAndAlign->second, AA);
415      coro::replaceCoroFree(CoroId, /*Elide=*/true);
416      NumOfCoroElided++;
417#ifndef NDEBUG
418      if (!CoroElideInfoOutputFilename.empty())
419        *getOrCreateLogFile()
420            << "Elide " << CoroId->getCoroutine()->getName() << " in "
421            << CoroId->getFunction()->getName() << "\n";
422#endif
423      ORE.emit([&]() {
424        return OptimizationRemark(DEBUG_TYPE, "CoroElide", CoroId)
425               << "'" << ore::NV("callee", CoroId->getCoroutine()->getName())
426               << "' elided in '"
427               << ore::NV("caller", CoroId->getFunction()->getName())
428               << "' (frame_size="
429               << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align="
430               << ore::NV("align", FrameSizeAndAlign->second.value()) << ")";
431      });
432    } else {
433      ORE.emit([&]() {
434        return OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId)
435               << "'" << ore::NV("callee", CoroId->getCoroutine()->getName())
436               << "' not elided in '"
437               << ore::NV("caller", CoroId->getFunction()->getName())
438               << "' (frame_size=unknown, align=unknown)";
439      });
440    }
441  }
442
443  return true;
444}
445
446static bool declaresCoroElideIntrinsics(Module &M) {
447  return coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.id.async"});
448}
449
450PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) {
451  auto &M = *F.getParent();
452  if (!declaresCoroElideIntrinsics(M))
453    return PreservedAnalyses::all();
454
455  Lowerer L(M);
456  L.CoroIds.clear();
457  L.collectPostSplitCoroIds(&F);
458  // If we did not find any coro.id, there is nothing to do.
459  if (L.CoroIds.empty())
460    return PreservedAnalyses::all();
461
462  AAResults &AA = AM.getResult<AAManager>(F);
463  DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
464  auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
465
466  bool Changed = false;
467  for (auto *CII : L.CoroIds)
468    Changed |= L.processCoroId(CII, AA, DT, ORE);
469
470  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
471}
472