[mlir] Delay adding the __resume function

[csigg]

Summary

The __resume function trips up LLVM's 'X86 DAG->DAG Instruction Selection' unless optimizations are disabled.

Only adding the __resume function when it's needed allows lowering through AsyncToLLVM and LLVM without '-O0' as long as the coroutine functionality is not used.

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 //===----------------------------------------------------------------------===//
 
 static constexpr const char *kResume = "__resume";
 
 /// A function that takes a coroutine handle and calls a `llvm.coro.resume`
 /// intrinsics. We need this function to be able to pass it to the async
 /// runtime execute API.
 static void addResumeFunction(ModuleOp module) {
+  if (module.lookupSymbol(kResume))
+    return;
+
   MLIRContext *ctx = module.getContext();
 
   OpBuilder moduleBuilder(module.getBody()->getTerminator());
   Location loc = module.getLoc();
 
-  if (module.lookupSymbol(kResume))
-    return;
-
   auto voidTy = LLVM::LLVMVoidType::get(ctx);
   auto i8Ptr = LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8));
 
   auto resumeOp = moduleBuilder.create<LLVM::LLVMFuncOp>(
       loc, kResume, LLVM::LLVMFunctionType::get(voidTy, {i8Ptr}));
   resumeOp.setPrivate();
 
   auto *block = resumeOp.addEntryBlock();
else if (operandType.isa<GroupType>())
       apiFuncName = kAwaitAllAndExecute;
     else
       return rewriter.notifyMatchFailure(op, "unsupported async type");
 
     Value operand = RuntimeAwaitAndResumeOpAdaptor(operands).operand();
     Value handle = RuntimeAwaitAndResumeOpAdaptor(operands).handle();
 
     // A pointer to coroutine resume intrinsic wrapper.
+    addResumeFunction(op->getParentOfType<ModuleOp>());
     auto resumeFnTy = AsyncAPI::resumeFunctionType(op->getContext());
     auto resumePtr = rewriter.create<LLVM::AddressOfOp>(
         op->getLoc(), LLVM::LLVMPointerType::get(resumeFnTy), kResume);
 
     rewriter.create<CallOp>(op->getLoc(), apiFuncName, TypeRange(),
                             ValueRange({operand, handle, resumePtr.res()}));
     rewriter.eraseOp(op);
 

 class RuntimeResumeOpLowering : public OpConversionPattern<RuntimeResumeOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeResumeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     // A pointer to coroutine resume intrinsic wrapper.
+    addResumeFunction(op->getParentOfType<ModuleOp>());
     auto resumeFnTy = AsyncAPI::resumeFunctionType(op->getContext());
     auto resumePtr = rewriter.create<LLVM::AddressOfOp>(
         op->getLoc(), LLVM::LLVMPointerType::get(resumeFnTy), kResume);
 
     // Call async runtime API to execute a coroutine in the managed thread.
     auto coroHdl = RuntimeResumeOpAdaptor(operands).handle();
     rewriter.replaceOpWithNewOp<CallOp>(op, TypeRange(), kExecute,
                                         ValueRange({coroHdl, resumePtr.res()}));
struct ConvertAsyncToLLVMPass
   void runOnOperation() override;
 };
 } // namespace
 
 void ConvertAsyncToLLVMPass::runOnOperation() {
   ModuleOp module = getOperation();
   MLIRContext *ctx = module->getContext();
 
-  // Add declarations for all functions required by the coroutines lowering.
-  addResumeFunction(module);
+  // Add declarations for most functions required by the coroutines lowering.
+  // We delay adding the resume function until it's needed because it currently
+  // fails to compile unless '-O0' is specified.
   addAsyncRuntimeApiDeclarations(module);
   addCRuntimeDeclarations(module);
 
   // Lower async.runtime and async.coro operations to Async Runtime API and
   // LLVM coroutine intrinsics.
 
   // Convert async dialect types and operations to LLVM dialect.
   AsyncRuntimeTypeConverter converter;