Merge pull request #34793 from aschwaighofer/irgen_async_silgen_name

IRGen: Allow calling a pre-defined async silgen_name function using the async convention

Author: DougGregor

lib/IRGen/Callee.h

@@ -155,23 +155,28 @@ namespace irgen {
 
     Signature Sig;
 
+    bool isFunctionPointerWithoutContext = false;
+
   public:
     /// Construct a FunctionPointer for an arbitrary pointer value.
     /// We may add more arguments to this; try to use the other
     /// constructors/factories if possible.
     explicit FunctionPointer(KindTy kind, llvm::Value *value,
                              PointerAuthInfo authInfo,
-                             const Signature &signature)
-        : Kind(kind), Value(value), AuthInfo(authInfo), Sig(signature) {
+                             const Signature &signature,
+                             bool isWithoutCtxt = false)
+        : Kind(kind), Value(value), AuthInfo(authInfo), Sig(signature),
+          isFunctionPointerWithoutContext(isWithoutCtxt) {
       // The function pointer should have function type.
       assert(value->getType()->getPointerElementType()->isFunctionTy());
       // TODO: maybe assert similarity to signature.getType()?
     }
 
     // Temporary only!
     explicit FunctionPointer(KindTy kind, llvm::Value *value,
-                             const Signature &signature)
-        : FunctionPointer(kind, value, PointerAuthInfo(), signature) {}
+                             const Signature &signature, bool
+                             isWithoutCtxt = false)
+        : FunctionPointer(kind, value, PointerAuthInfo(), signature, isWithoutCtxt) {}
 
     static FunctionPointer forDirect(IRGenModule &IGM,
                                      llvm::Constant *value,
@@ -243,6 +248,10 @@ namespace irgen {
 
     /// Form a FunctionPointer whose KindTy is ::Function.
     FunctionPointer getAsFunction(IRGenFunction &IGF) const;
+
+    bool useStaticContextSize() const {
+      return isFunctionPointerWithoutContext;
+    }
   };
 
   class Callee {

lib/IRGen/GenCall.cpp

@@ -1830,7 +1830,7 @@ void irgen::extractScalarResults(IRGenFunction &IGF, llvm::Type *bodyType,
 std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
     IRGenFunction &IGF, SILFunctionTypeRepresentation representation,
     FunctionPointer functionPointer, llvm::Value *thickContext,
-    std::pair<bool, bool> values) {
+    std::pair<bool, bool> values, Size initialContextSize) {
   assert(values.first || values.second);
   bool emitFunction = values.first;
   bool emitSize = values.second;
@@ -1997,11 +1997,16 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
     }
     llvm::Value *size = nullptr;
     if (emitSize) {
-      auto *ptr = functionPointer.getRawPointer();
-      auto *descriptorPtr =
-          IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
-      auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
-      size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
+      if (functionPointer.useStaticContextSize()) {
+        size = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
+                                      initialContextSize.getValue());
+      }  else {
+        auto *ptr = functionPointer.getRawPointer();
+        auto *descriptorPtr =
+            IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
+        auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
+        size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
+      }
     }
     return {fn, size};
   }
@@ -2270,7 +2275,8 @@ class AsyncCallEmission final : public CallEmission {
     llvm::Value *dynamicContextSize32;
     std::tie(calleeFunction, dynamicContextSize32) = getAsyncFunctionAndSize(
         IGF, CurCallee.getOrigFunctionType()->getRepresentation(),
-        CurCallee.getFunctionPointer(), thickContext);
+        CurCallee.getFunctionPointer(), thickContext,
+        std::make_pair(true, true), layout.getSize());
     auto *dynamicContextSize =
         IGF.Builder.CreateZExt(dynamicContextSize32, IGF.IGM.SizeTy);
     contextBuffer = emitAllocAsyncContext(IGF, dynamicContextSize);
@@ -4494,13 +4500,20 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
   switch (Kind.value) {
   case KindTy::Value::Function:
     return Value;
-  case KindTy::Value::AsyncFunctionPointer:
-    auto *descriptorPtr =
-        IGF.Builder.CreateBitCast(Value, IGF.IGM.AsyncFunctionPointerPtrTy);
-    auto *addrPtr = IGF.Builder.CreateStructGEP(descriptorPtr, 0);
-    return IGF.emitLoadOfRelativePointer(
-        Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
-        /*expectedType*/ getFunctionType()->getPointerTo());
+  case KindTy::Value::AsyncFunctionPointer: {
+    if (!isFunctionPointerWithoutContext) {
+      auto *descriptorPtr =
+          IGF.Builder.CreateBitCast(Value, IGF.IGM.AsyncFunctionPointerPtrTy);
+      auto *addrPtr = IGF.Builder.CreateStructGEP(descriptorPtr, 0);
+      return IGF.emitLoadOfRelativePointer(
+          Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
+          /*expectedType*/ getFunctionType()->getPointerTo());
+    } else {
+      return IGF.Builder.CreateBitOrPointerCast(
+          Value, getFunctionType()->getPointerTo());
+    }
+  }
+
   }
 }
 

lib/IRGen/GenCall.h

@@ -313,7 +313,8 @@ namespace irgen {
   std::pair<llvm::Value *, llvm::Value *> getAsyncFunctionAndSize(
       IRGenFunction &IGF, SILFunctionTypeRepresentation representation,
       FunctionPointer functionPointer, llvm::Value *thickContext,
-      std::pair<bool, bool> values = {true, true});
+      std::pair<bool, bool> values = {true, true},
+      Size initialContextSize = Size(0));
   llvm::CallingConv::ID expandCallingConv(IRGenModule &IGM,
                                      SILFunctionTypeRepresentation convention);
 

lib/IRGen/IRGenSIL.cpp

@@ -2288,13 +2288,16 @@ void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
       fn, NotForDefinition, false /*isDynamicallyReplaceableImplementation*/,
       isa<PreviousDynamicFunctionRefInst>(i));
   llvm::Value *value;
-  if (fn->isAsync()) {
+  auto isSpecialAsyncWithoutCtxtSize =
+      fn->isAsync() && fn->getName().equals("swift_task_future_wait");
+  if (fn->isAsync() && !isSpecialAsyncWithoutCtxtSize) {
     value = IGM.getAddrOfAsyncFunctionPointer(fn);
     value = Builder.CreateBitCast(value, fnPtr->getType());
   } else {
     value = fnPtr;
   }
-  FunctionPointer fp = FunctionPointer(fnType, value, sig);
+  FunctionPointer fp =
+      FunctionPointer(fnType, value, sig, isSpecialAsyncWithoutCtxtSize);
 
   // Store the function as a FunctionPointer so we can avoid bitcasting
   // or thunking if we don't need to.

test/IRGen/async.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -primary-file %s -emit-ir -enable-experimental-concurrency | %FileCheck %s
+// RUN: %target-swift-frontend -primary-file %s -emit-ir -enable-experimental-concurrency | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
 
 // REQUIRES: concurrency
 
@@ -8,3 +8,19 @@ public func f() async { }
 // CHECK: "$s5async1gyyYKF"
 public func g() async throws { }
 
+
+public class SomeClass {}
+
+@_silgen_name("swift_task_future_wait")
+public func task_future_wait(_ task: __owned SomeClass) async throws -> Int
+
+// CHECK: define{{.*}} swiftcc void @"$s5async8testThisyyAA9SomeClassCnYF"(%swift.task* %0, %swift.executor* %1, %swift.context* %2)
+// CHECK-64: call swiftcc i8* @swift_task_alloc(%swift.task* %0, i64 64)
+// CHECK: tail call swiftcc void @swift_task_future_wait(
+public func testThis(_ task: __owned SomeClass) async {
+  do {
+    let _ = try await task_future_wait(task)
+  } catch _ {
+    print("error")
+  }
+}