Implement LLVM bits

Author: yuxuanchen1997

llvm/include/llvm/IR/Instruction.h

@@ -445,6 +445,10 @@ class Instruction : public User,
   /// !annotation metadata, append the tuple to
   /// the existing node.
   void addAnnotationMetadata(SmallVector<StringRef> Annotations);
+
+  /// Returns true if an !annotation metadata is set to this instruction.
+  bool hasAnnotationMetadata(StringRef Name) const;
+
   /// Returns the AA metadata for this instruction.
   AAMDNodes getAAMetadata() const;
 

llvm/include/llvm/Transforms/Coroutines/CoroAnnotationElide.h

@@ -0,0 +1,34 @@
+//===- CoroAnnotationElide.h - Optimizing a coro_must_elide call ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// \file
+// This pass transforms all Call or Invoke instructions that are annotated
+// "coro_must_elide" to call the `.noalloc` variant of coroutine instead.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
+#define LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
+
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LazyCallGraph.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+struct CoroAnnotationElidePass : PassInfoMixin<CoroAnnotationElidePass> {
+  CoroAnnotationElidePass() {}
+
+  PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
+                        LazyCallGraph &CG, CGSCCUpdateResult &UR);
+
+  static bool isRequired() { return false; }
+};
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H

llvm/lib/IR/Metadata.cpp

@@ -1703,6 +1703,22 @@ void Instruction::addAnnotationMetadata(StringRef Name) {
   setMetadata(LLVMContext::MD_annotation, MD);
 }
 
+bool Instruction::hasAnnotationMetadata(StringRef Name) const {
+  auto *Metadata = getMetadata(LLVMContext::MD_annotation);
+  if (!Metadata)
+    return false;
+
+  auto *Tuple = cast<MDTuple>(Metadata);
+  for (auto &N : Tuple->operands()) {
+    if (auto *S = dyn_cast<MDString>(N.get())) {
+      if (S->getString() == Name) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 AAMDNodes Instruction::getAAMetadata() const {
   AAMDNodes Result;
   // Not using Instruction::hasMetadata() because we're not interested in

llvm/lib/Passes/PassBuilder.cpp

@@ -134,6 +134,7 @@
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
 #include "llvm/Transforms/CFGuard.h"
+#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
 #include "llvm/Transforms/Coroutines/CoroConditionalWrapper.h"
 #include "llvm/Transforms/Coroutines/CoroEarly.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -32,6 +32,7 @@
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
+#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
 #include "llvm/Transforms/Coroutines/CoroConditionalWrapper.h"
 #include "llvm/Transforms/Coroutines/CoroEarly.h"
@@ -968,8 +969,8 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
   // it's been modified since.
   MainCGPipeline.addPass(createCGSCCToFunctionPassAdaptor(
       RequireAnalysisPass<ShouldNotRunFunctionPassesAnalysis, Function>()));
-
   MainCGPipeline.addPass(CoroSplitPass(Level != OptimizationLevel::O0));
+  MainCGPipeline.addPass(CoroAnnotationElidePass());
 
   // Make sure we don't affect potential future NoRerun CGSCC adaptors.
   MIWP.addLateModulePass(createModuleToFunctionPassAdaptor(

llvm/lib/Passes/PassRegistry.def

@@ -234,6 +234,7 @@ CGSCC_PASS("attributor-light-cgscc", AttributorLightCGSCCPass())
 CGSCC_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 CGSCC_PASS("no-op-cgscc", NoOpCGSCCPass())
 CGSCC_PASS("openmp-opt-cgscc", OpenMPOptCGSCCPass())
+CGSCC_PASS("coro-annotation-elide", CoroAnnotationElidePass())
 #undef CGSCC_PASS
 
 #ifndef CGSCC_PASS_WITH_PARAMS

llvm/lib/Transforms/Coroutines/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_llvm_component_library(LLVMCoroutines
   Coroutines.cpp
+  CoroAnnotationElide.cpp
   CoroCleanup.cpp
   CoroConditionalWrapper.cpp
   CoroEarly.cpp

llvm/lib/Transforms/Coroutines/CoroAnnotationElide.cpp

@@ -0,0 +1,117 @@
+//===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
+
+#include "llvm/Analysis/LazyCallGraph.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+
+#include <cassert>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "coro-annotation-elide"
+
+#define CORO_MUST_ELIDE_ANNOTATION "coro_must_elide"
+
+static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
+  for (Instruction &I : F->getEntryBlock())
+    if (!isa<AllocaInst>(&I))
+      return &I;
+  llvm_unreachable("no terminator in the entry block");
+}
+
+static Value *allocateFrameInCaller(Function *Caller, uint64_t FrameSize,
+                                    Align FrameAlign) {
+  LLVMContext &C = Caller->getContext();
+  BasicBlock::iterator InsertPt =
+      getFirstNonAllocaInTheEntryBlock(Caller)->getIterator();
+  const DataLayout &DL = Caller->getDataLayout();
+  auto FrameTy = ArrayType::get(Type::getInt8Ty(C), FrameSize);
+  auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);
+  Frame->setAlignment(FrameAlign);
+  return new BitCastInst(Frame, PointerType::getUnqual(C), "vFrame", InsertPt);
+}
+
+static void processCall(CallBase *CB, Function *Caller, Function *NewCallee,
+                        uint64_t FrameSize, Align FrameAlign) {
+  auto *FramePtr = allocateFrameInCaller(Caller, FrameSize, FrameAlign);
+  CB->setCalledFunction(NewCallee->getFunctionType(), NewCallee);
+  auto NewCBInsertPt = CB->getIterator();
+  llvm::CallBase *NewCB = nullptr;
+  SmallVector<Value *, 4> NewArgs;
+  NewArgs.append(CB->arg_begin(), CB->arg_end());
+  NewArgs.push_back(FramePtr);
+
+  // TODO: See CallBase::Create(CallBase*, ...)
+  if (auto *CI = dyn_cast<CallInst>(CB)) {
+    auto *NewCI = CallInst::Create(NewCallee->getFunctionType(), NewCallee,
+                                   NewArgs, "", NewCBInsertPt);
+    NewCI->setTailCallKind(CI->getTailCallKind());
+    NewCB = NewCI;
+  } else if (auto *II = dyn_cast<InvokeInst>(CB)) {
+    NewCB = InvokeInst::Create(NewCallee->getFunctionType(), NewCallee,
+                               II->getNormalDest(), II->getUnwindDest(),
+                               NewArgs, std::nullopt, "", NewCBInsertPt);
+  } else {
+    llvm_unreachable("CallBase should either be Call or Invoke!");
+  }
+
+  NewCB->setCallingConv(CB->getCallingConv());
+  NewCB->setAttributes(CB->getAttributes());
+  NewCB->setDebugLoc(CB->getDebugLoc());
+  std::copy(CB->bundle_op_info_begin(), CB->bundle_op_info_end(),
+            NewCB->bundle_op_info_begin());
+
+  CB->replaceAllUsesWith(NewCB);
+  CB->eraseFromParent();
+}
+
+PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
+                                               CGSCCAnalysisManager &AM,
+                                               LazyCallGraph &CG,
+                                               CGSCCUpdateResult &UR) {
+  bool Changed = false;
+
+  for (LazyCallGraph::Node &N : C) {
+    Function *Callee = &N.getFunction();
+    Function *NewCallee = Callee->getParent()->getFunction(
+        (Callee->getName() + ".noalloc").str());
+    if (!NewCallee) {
+      continue;
+    }
+
+    auto FrameSize = NewCallee->getParamDereferenceableBytes(0);
+    auto FrameAlign = NewCallee->getParamAlign(0).valueOrOne();
+    llvm::dbgs() << "Callee: ";
+    Callee->dump();
+
+    SmallVector<CallBase *, 4> Users;
+    for (auto *U : Callee->users()) {
+      if (auto *CB = dyn_cast<CallBase>(U)) {
+        Users.push_back(CB);
+      }
+    }
+
+    for (auto *CB : Users) {
+      auto *Caller = CB->getFunction();
+      if (Caller && Caller->isPresplitCoroutine() &&
+          CB->hasAnnotationMetadata(CORO_MUST_ELIDE_ANNOTATION)) {
+        processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
+        Changed = true;
+      }
+    }
+  }
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}

llvm/lib/Transforms/Coroutines/CoroInternal.h

@@ -26,6 +26,10 @@ bool declaresIntrinsics(const Module &M,
                         const std::initializer_list<StringRef>);
 void replaceCoroFree(CoroIdInst *CoroId, bool Elide);
 
+void suppressCoroAllocs(CoroIdInst *CoroId);
+void suppressCoroAllocs(LLVMContext &Context,
+                        ArrayRef<CoroAllocInst *> CoroAllocs);
+
 /// Attempts to rewrite the location operand of debug intrinsics in terms of
 /// the coroutine frame pointer, folding pointer offsets into the DIExpression
 /// of the intrinsic.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -25,6 +25,7 @@
 #include "llvm/ADT/PriorityWorklist.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/CFG.h"
@@ -1179,6 +1180,14 @@ static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
   Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);
 }
 
+static TypeSize getFrameSizeForShape(coro::Shape &Shape) {
+  // In the same function all coro.sizes should have the same result type.
+  auto *SizeIntrin = Shape.CoroSizes.back();
+  Module *M = SizeIntrin->getModule();
+  const DataLayout &DL = M->getDataLayout();
+  return DL.getTypeAllocSize(Shape.FrameTy);
+}
+
 static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
   if (Shape.ABI == coro::ABI::Async)
     updateAsyncFuncPointerContextSize(Shape);
@@ -1194,10 +1203,8 @@ static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
 
   // In the same function all coro.sizes should have the same result type.
   auto *SizeIntrin = Shape.CoroSizes.back();
-  Module *M = SizeIntrin->getModule();
-  const DataLayout &DL = M->getDataLayout();
-  auto Size = DL.getTypeAllocSize(Shape.FrameTy);
-  auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
+  auto *SizeConstant =
+      ConstantInt::get(SizeIntrin->getType(), getFrameSizeForShape(Shape));
 
   for (CoroSizeInst *CS : Shape.CoroSizes) {
     CS->replaceAllUsesWith(SizeConstant);
@@ -2077,6 +2084,46 @@ static void addPrepareFunction(const Module &M,
     Fns.push_back(PrepareFn);
 }
 
+static Function *createNoAllocVariant(Function &F, ValueToValueMapTy &VMap,
+                                      coro::Shape &Shape) {
+  auto *OrigFnTy = F.getFunctionType();
+  auto OldParams = OrigFnTy->params();
+
+  SmallVector<Type *> NewParams;
+  NewParams.reserve(OldParams.size() + 1);
+  for (Type *T : OldParams) {
+    NewParams.push_back(T);
+  }
+  NewParams.push_back(PointerType::getUnqual(Shape.FrameTy));
+
+  auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,
+                                    OrigFnTy->isVarArg());
+  Function *NoAllocF =
+      Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");
+  unsigned int Idx = 0;
+  for (const auto &I : F.args()) {
+    VMap[&I] = NoAllocF->getArg(Idx++);
+  }
+  SmallVector<ReturnInst *, 4> Returns;
+  CloneFunctionInto(NoAllocF, &F, VMap,
+                    CloneFunctionChangeType::LocalChangesOnly, Returns);
+
+  if (Shape.CoroBegin) {
+    auto *NewCoroBegin = cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);
+    auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());
+    coro::replaceCoroFree(NewCoroId, /*Elide=*/true);
+    coro::suppressCoroAllocs(NewCoroId);
+    NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(Idx));
+    NewCoroBegin->eraseFromParent();
+  }
+
+  Module *M = F.getParent();
+  M->getFunctionList().insert(M->end(), NoAllocF);
+
+  removeUnreachableBlocks(*NoAllocF);
+  return NoAllocF;
+}
+
 CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
     : MaterializableCallback(coro::defaultMaterializable),
       OptimizeFrame(OptimizeFrame) {}
@@ -2108,13 +2155,14 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
+
     LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
                       << "\n");
     F.setSplittedCoroutine();
 
     SmallVector<Function *, 4> Clones;
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
-    const coro::Shape Shape =
+    coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
     removeCoroEndsFromRampFunction(Shape);
@@ -2133,11 +2181,28 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
       for (Function *Clone : Clones)
         UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
     }
-  }
 
-    for (auto *PrepareFn : PrepareFns) {
-      replaceAllPrepares(PrepareFn, CG, C);
+    if (Shape.ABI == coro::ABI::Switch) {
+      ValueToValueMapTy VMap;
+      auto *NoAllocF = createNoAllocVariant(F, VMap, Shape);
+      NoAllocF->addFnAttr("elided-coro");
+      auto NewAttrs = NoAllocF->getAttributes();
+      // We just appended the frame pointer as the last argument of the new
+      // function.
+      auto FrameIdx = NoAllocF->arg_size() - 1;
+      // When we elide allocation, we read these attributes to determine the
+      // frame size and alignment.
+      addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,
+                           Shape.FrameSize, Shape.FrameAlign,
+                           /*NoAlias=*/false);
+
+      NoAllocF->setAttributes(NewAttrs);
     }
+  }
+
+  for (auto *PrepareFn : PrepareFns) {
+    replaceAllPrepares(PrepareFn, CG, C);
+  }
 
   return PreservedAnalyses::none();
 }