[CoroElide] Remove fallback for frame layout determination

[nikic]

Summary

Only determine the frame layout based on dereferenceable and align attributes, and remove the type-based fallback, which is incompatible with opaque pointers. The dereferenceable attribute is required, while the align attribute uses default alignment of 1 (commonly, align 1 attributes do not get placed, relying on default alignment).

The CoroSplit pass producing the resume function adds the necessary attributes in https://github.com/llvm/llvm-project/blob/7daed359111f6d151fef447f520f85ef1dabedf6/llvm/lib/Transforms/Coroutines/CoroSplit.cpp#L840, and they are checked in coro-debug.ll.

[ChuanqiXu]

nit in the summary: the dereferenceable attribute is added in middle-end since *.resume function is generated in middle-end (CoroSplit pass). Maybe we need a test for it (check dereferenceable attribute for *.resume function)

[nikic]

In D120988#3362716, @ChuanqiXu wrote:

nit in the summary: the dereferenceable attribute is added in middle-end since *.resume function is generated in middle-end (CoroSplit pass). Maybe we need a test for it (check dereferenceable attribute for *.resume function)

Thanks for pointing that out, I've updated the description. Apparently there already is a test for the emitted attributes in coro-debug.ll and a couple other test files as well.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

for (Instruction &I : instructions(F))
     if (auto *Call = dyn_cast<CallInst>(&I))
       if (Call->isTailCall() && operandReferences(Call, Frame, AA) &&
           !Call->isMustTailCall())
         Call->setTailCall(false);
 }
 
 // Given a resume function @f.resume(%f.frame* %frame), returns the size
 // and expected alignment of %f.frame type.
-static std::pair<uint64_t, Align> getFrameLayout(Function *Resume) {
-  // Prefer to pull information from the function attributes.
+static Optional<std::pair<uint64_t, Align>> getFrameLayout(Function *Resume) {
+  // Pull information from the function attributes.
   auto Size = Resume->getParamDereferenceableBytes(0);
-  auto Align = Resume->getParamAlign(0);
-
-  // If those aren't given, extract them from the type.
-  if (Size == 0 || !Align) {
-    auto *FrameTy = Resume->arg_begin()->getType()->getPointerElementType();
-
-    const DataLayout &DL = Resume->getParent()->getDataLayout();
-    if (!Size) Size = DL.getTypeAllocSize(FrameTy);
-    if (!Align) Align = DL.getABITypeAlign(FrameTy);
-  }
-
-  return std::make_pair(Size, *Align);
+  if (!Size)
+    return None;
+  return std::make_pair(Size, Resume->getParamAlign(0).valueOrOne());
 }
 
 // Finds first non alloca instruction in the entry block of a function.
 static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
   for (Instruction &I : F->getEntryBlock())
     if (!isa<AllocaInst>(&I))
       return &I;
   llvm_unreachable("no terminator in the entry block");
bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
   auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
       Resumers,
       ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
 
   for (auto &It : DestroyAddr)
     replaceWithConstant(DestroyAddrConstant, It.second);
 
   if (ShouldElide) {
-    auto FrameSizeAndAlign = getFrameLayout(cast<Function>(ResumeAddrConstant));
-    elideHeapAllocations(CoroId->getFunction(), FrameSizeAndAlign.first,
-                         FrameSizeAndAlign.second, AA);
+    if (auto FrameSizeAndAlign =
+            getFrameLayout(cast<Function>(ResumeAddrConstant))) {
+      elideHeapAllocations(CoroId->getFunction(), FrameSizeAndAlign->first,
+                           FrameSizeAndAlign->second, AA);
       coro::replaceCoroFree(CoroId, /*Elide=*/true);
       NumOfCoroElided++;
 #ifndef NDEBUG
       if (!CoroElideInfoOutputFilename.empty())
         *getOrCreateLogFile()
             << "Elide " << CoroId->getCoroutine()->getName() << " in "
             << CoroId->getFunction()->getName() << "\n";
 #endif
     }
+  }
 
   return true;
 }
 
 // See if there are any coro.subfn.addr instructions referring to coro.devirt
 // trigger, if so, replace them with a direct call to devirt trigger function.
 static bool replaceDevirtTrigger(Function &F) {
   SmallVector<CoroSubFnInst *, 1> DevirtAddr;

llvm/test/Transforms/Coroutines/coro-elide-musttail.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; Only run with new pass manager since old pass manager's alias analysis isn't
 ; powerful enough to tell that the tailcall's arguments don't alias the frame.
 ;
 ; RUN: opt < %s -passes='coro-elide' -S | FileCheck %s
 ; RUN: opt < %s -aa-pipeline= -passes='coro-elide' -S | FileCheck %s -check-prefix=NOAA
 
 %"bar.Frame" = type { void (%"bar.Frame"*)*, void (%"bar.Frame"*)*, %"struct.coroutine<false, int>::promise_type", i1 }
 %"struct.coroutine<false, int>::promise_type" = type { i32 }
 %"foo.Frame" = type { void (%"foo.Frame"*)*, void (%"foo.Frame"*)*, %"struct.coroutine<true, int>::promise_type", i1, %"struct.coroutine<true, int>::promise_type::final_awaitable" }
 %"struct.coroutine<true, int>::promise_type" = type { i32 }
 %"struct.coroutine<true, int>::promise_type::final_awaitable" = type { i8 }
 @"bar.resumers" = private constant [3 x void (%"bar.Frame"*)*] [void (%"bar.Frame"*)* @"bar.resume", void (%"bar.Frame"*)* undef, void (%"bar.Frame"*)* undef]
 
 declare dso_local void @"bar"() align 2
-declare dso_local fastcc void @"bar.resume"(%"bar.Frame"*) align 2
+declare dso_local fastcc void @"bar.resume"(%"bar.Frame"* align 8 dereferenceable(24)) align 2
 
 ; There is a musttail call.
 ; With alias analysis, we can tell that the frame does not interfere with CALL34, and hence we can keep the tailcalls.
 ; Without alias analysis, we have to keep the tailcalls.
 define internal fastcc void @foo.resume_musttail(%"foo.Frame"* %FramePtr) {
 ; CHECK-LABEL: @foo.resume_musttail(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = alloca [24 x i8], align 8

llvm/test/Transforms/Coroutines/coro-elide-stat.ll

 
 ; CHECK: 2 coro-elide  - The # of coroutine get elided.
 ; FILE: Elide f in callResume
 ; FILE: Elide f in callResumeMultiRetDommmed
 
 declare void @print(i32) nounwind
 
 ; resume part of the coroutine
-define fastcc void @f.resume(i8*) {
+define fastcc void @f.resume(i8* dereferenceable(1)) {
   tail call void @print(i32 0)
   ret void
 }
 
 ; destroy part of the coroutine
 define fastcc void @f.destroy(i8*) {
   tail call void @print(i32 1)
   ret void

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

 ; Tests that the coro.destroy and coro.resume are devirtualized where possible,
 ; SCC pipeline restarts and inlines the direct calls.
 ; RUN: opt < %s -S \
 ; RUN:   -passes='cgscc(repeat<2>(inline,function(coro-elide,dce)))' \
 ; RUN:   | FileCheck %s
 
 declare void @print(i32) nounwind
 
 ; resume part of the coroutine
-define fastcc void @f.resume(i8*) {
+define fastcc void @f.resume(i8* dereferenceable(1)) {
   tail call void @print(i32 0)
   ret void
 }
 
 ; destroy part of the coroutine
 define fastcc void @f.destroy(i8*) {
   tail call void @print(i32 1)
   ret void

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

 ; Tests that the dynamic allocation and deallocation of the coroutine frame is
 ; elided and any tail calls referencing the coroutine frame has the tail 
 ; call attribute removed.
 ; RUN: opt < %s -S \
 ; RUN:   -passes='cgscc(inline,function(coro-elide,instsimplify,simplifycfg))' \
 ; RUN:   -aa-pipeline='basic-aa' | FileCheck %s
 
 declare void @print(i32) nounwind
 
 %f.frame = type {i32}
 
 declare void @bar(i8*)
 
-declare fastcc void @f.resume(%f.frame*)
+declare fastcc void @f.resume(%f.frame* align 4 dereferenceable(4))
 declare fastcc void @f.destroy(%f.frame*)
 declare fastcc void @f.cleanup(%f.frame*)
 
 declare void @may_throw()
 declare i8* @CustomAlloc(i32)
 declare void @CustomFree(i8*)
 
 @f.resumers = internal constant [3 x void (%f.frame*)*]