[Coroutines] PR34897: Fix incorrect elisions

[modocache]

Summary

https://bugs.llvm.org/show_bug.cgi?id=34897 demonstrates an incorrect
coroutine frame allocation elision in the coro-elide pass. The elision
is performed on the basis that the SSA variables from all llvm.coro.begin
are directly referenced in subsequent llvm.coro.destroy instructions.

However, this ignores the fact that the function may exit through paths
that do not run these destroy instructions. In the sample program from
PR34897, for example, the llvm.coro.destroy instruction is only
executed in exception handling code. When the coroutine function exits
normally, llvm.coro.destroy is not called. Eliding the allocation in
this case causes a subsequent reference to the coroutine handle from
outside of the function to access freed memory.

To fix the issue, when finding an llvm.coro.destroy for each llvm.coro.begin,
only consider llvm.coro.destroy that are executed along non-exceptional paths.

Test Plan:

1. Download the sample program from https://bugs.llvm.org/show_bug.cgi?id=34897, compile it with clang++ -fcoroutines-ts -stdlib=libc++ -std=c++1z -O2, and run it. It should print "run1\ncheck1\nrun2\ncheck2" and then exit successfully.
2. Compile https://godbolt.org/g/mCKfnr and confirm it is still optimized to a single instruction, 'return 1190'.
3. check-llvm

[modocache]

Friendly ping! Let me know if this seems totally wrong to you, @GorNishanov :)

[modocache]

Just tidying up one line with clang-format. I guess everyone's busy preparing for the Jacksonville standards meeting. Please give this a review when you all have a second -- in addition to the reported bug, https://bugs.llvm.org/show_bug.cgi?id=34897, I've also confirmed that this diff fixes ASAN errors that surfaced in a project at work. Thanks! :)

[GorNishanov]

I believe that this check is too aggressive. It prevents http://godbolt.org/g/26viuZ from optimizing as well.
I think a distinction should be made for exceptional paths and happy path.
We need to make sure that we call coro.destroy on all happy paths.

I did not page-in everything I thought about the when_all_test bug, but, I scribbled somewhere this note about it: "Ignore coro.destroy on exceptional paths for the purpose of heap elision".

[modocache]

Sorry it took me a while to get around to this, but I think this new approach fits more along the lines of what you were thinking, @GorNishanov. I also confirmed that it optimizes the program you linked to just as well as before.

[GorNishanov]

LGTM!

[modocache]

Hooray! Thanks again for the review, @GorNishanov.

lib/Transforms/Coroutines/CoroElide.cpp

 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass replaces dynamic allocation of coroutine frame with alloca and
 // replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls
 // to coroutine sub-functions.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-elide"
 
 namespace {
 // Created on demand if CoroElide pass has work to do.
 struct Lowerer : coro::LowererBase {
   SmallVector<CoroIdInst *, 4> CoroIds;
   SmallVector<CoroBeginInst *, 1> CoroBegins;
   SmallVector<CoroAllocInst *, 1> CoroAllocs;
   SmallVector<CoroSubFnInst *, 4> ResumeAddr;
   SmallVector<CoroSubFnInst *, 4> DestroyAddr;
   SmallVector<CoroFreeInst *, 1> CoroFrees;
 
   Lowerer(Module &M) : LowererBase(M) {}
 
   void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
-  bool shouldElide() const;
-  bool processCoroId(CoroIdInst *, AAResults &AA);
+  bool shouldElide(Function *F, DominatorTree &DT) const;
+  bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT);
 };
 } // end anonymous namespace
 
 // Go through the list of coro.subfn.addr intrinsics and replace them with the
 // provided constant.
 static void replaceWithConstant(Constant *Value,
                                 SmallVectorImpl<CoroSubFnInst *> &Users) {
   if (Users.empty())
for (auto *CB : CoroBegins) {
     CB->eraseFromParent();
   }
 
   // Since now coroutine frame lives on the stack we need to make sure that
   // any tail call referencing it, must be made non-tail call.
   removeTailCallAttribute(Frame, AA);
 }
 
-bool Lowerer::shouldElide() const {
+bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const {
   // If no CoroAllocs, we cannot suppress allocation, so elision is not
   // possible.
   if (CoroAllocs.empty())
     return false;
 
   // Check that for every coro.begin there is a coro.destroy directly
-  // referencing the SSA value of that coro.begin. If the value escaped, then
-  // coro.destroy would have been referencing a memory location storing that
-  // value and not the virtual register.
+  // referencing the SSA value of that coro.begin for every reachable exiting
+  // path out of the function. If the value escaped, then coro.destroy would
+  // have been referencing a memory location storing that value and not the
+  // virtual register.
+
+  SmallPtrSet<Instruction *, 8> TerminatorInsts;
+  for (BasicBlock &B : *F) {
+    auto *TI = B.getTerminator();
+    if (isa<ReturnInst>(TI) || isa<ResumeInst>(TI))
+      TerminatorInsts.insert(TI);
+  }
 
   SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
-
+  SmallPtrSet<Instruction *, 8> DominatedTerminatorInsts;
   for (CoroSubFnInst *DA : DestroyAddr) {
-    if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame()))
+    if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame())) {
       ReferencedCoroBegins.insert(CB);
-    else
+
+      // Update the set of terminator instructions that are dominated by one of
+      // these destroy instructions.
+      for (Instruction *I : TerminatorInsts) {
+        // No need to check if the function has already determined that all
+        // terminator instructions are dominated.
+        if (DominatedTerminatorInsts.size() == TerminatorInsts.size())
+          break;
+        // Skip instructions that we know are dominated.
+        if (DominatedTerminatorInsts.count(I) != 0)
+          continue;
+
+        if (DT.dominates(DA, I))
+          DominatedTerminatorInsts.insert(I);
+      }
+    } else
       return false;
   }
 
-  // If size of the set is the same as total number of CoroBegins, means we
-  // found a coro.free or coro.destroy mentioning a coro.begin and we can
-  // perform heap elision.
-  return ReferencedCoroBegins.size() == CoroBegins.size();
+  // If size of the set is the same as total number of CoroBegins, that means
+  // that we found a coro.free or coro.destroy that directly refences the
+  // coro.begin. If that set of coro.free or coro.destroy dominates each
+  // terminator instruction in the function, we can perform heap elision.
+  return DominatedTerminatorInsts.size() == TerminatorInsts.size() &&
+         ReferencedCoroBegins.size() == CoroBegins.size();
 }
 
-bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA) {
+bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
+                            DominatorTree &DT) {
   CoroBegins.clear();
   CoroAllocs.clear();
   CoroFrees.clear();
   ResumeAddr.clear();
   DestroyAddr.clear();
 
   // Collect all coro.begin and coro.allocs associated with this coro.id.
   for (User *U : CoroId->users()) {
bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
   ConstantArray *Resumers = CoroId->getInfo().Resumers;
   assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
                      "of coroutine subfunctions");
   auto *ResumeAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
 
   replaceWithConstant(ResumeAddrConstant, ResumeAddr);
 
-  bool ShouldElide = shouldElide();
+  bool ShouldElide = shouldElide(CoroId->getFunction(), DT);
 
   auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
       Resumers,
       ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
 
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
 
   if (ShouldElide) {
for (auto &I : instructions(F))
           if (CII->getCoroutine() != CII->getFunction())
             L->CoroIds.push_back(CII);
 
     // If we did not find any coro.id, there is nothing to do.
     if (L->CoroIds.empty())
       return Changed;
 
     AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+    DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
     for (auto *CII : L->CoroIds)
-      Changed |= L->processCoroId(CII, AA);
+      Changed |= L->processCoroId(CII, AA, DT);
 
     return Changed;
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AAResultsWrapperPass>();
+    AU.addRequired<DominatorTreeWrapperPass>();
   }
   StringRef getPassName() const override { return "Coroutine Elision"; }
 };
 }
 
 char CoroElide::ID = 0;
 INITIALIZE_PASS_BEGIN(
     CoroElide, "coro-elide",

test/Transforms/Coroutines/coro-heap-elide.ll

; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %3 = bitcast i8* %2 to void (i8*)*
   call fastcc void %3(i8* %hdl)
 
 ; CHECK-NEXT: ret void
   ret void
 }
 
+; CHECK-LABEL: @callResume_PR34897_no_elision(
+define void @callResume_PR34897_no_elision(i1 %cond) {
+; CHECK-LABEL: entry:
+entry:
+; CHECK: call i8* @CustomAlloc(
+  %hdl = call i8* @f()
+; CHECK: tail call void @bar(
+  tail call void @bar(i8* %hdl)
+; CHECK: tail call void @bar(
+  tail call void @bar(i8* null)
+  br i1 %cond, label %if.then, label %if.else
+
+; CHECK-LABEL: if.then:
+if.then:
+; CHECK: call fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8*
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  call fastcc void %1(i8* %hdl)
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.destroy to void (i8*)*)(i8*
+  %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %3 = bitcast i8* %2 to void (i8*)*
+  call fastcc void %3(i8* %hdl)
+  br label %return
+
+if.else:
+  br label %return
+
+; CHECK-LABEL: return:
+return:
+; CHECK: ret void
+  ret void
+}
+
+
 ; a coroutine start function (cannot elide heap alloc, due to second argument to
 ; coro.begin not pointint to coro.alloc)
 define i8* @f_no_elision() personality i8* null {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
                       i8* bitcast (i8*()* @f_no_elision to i8*),
                       i8* bitcast ([3 x void (%f.frame*)*]* @f.resumers to i8*))
   %alloc = call i8* @CustomAlloc(i32 4)