[Distributed] IRGen: Adjust distributed method accessor to store result into a provided buffer

xedin · xedin · commit 111ceb189a15 · 2021-12-17T10:52:56.000-08:00
Instead of trying to return result from distributed thunk directly,
modify accessor to store result into the caller-provided buffer.
Doing so helps us avoid boxing the result into `Any`.
diff --git a/lib/IRGen/GenDistributed.cpp b/lib/IRGen/GenDistributed.cpp
@@ -98,7 +98,7 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
                                           SILFunction *DistMethod) {
   auto &Context = IGM.Context;
 
-  auto getParamForArguments = [&]() {
+  auto getRawPointerParmeter = [&]() {
     auto ptrType = Context.getUnsafeRawPointerType();
     return SILParameterInfo(ptrType->getCanonicalType(),
                             ParameterConvention::Direct_Guaranteed);
@@ -121,9 +121,11 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
   // its result(s) out.
   return SILFunctionType::get(
       /*genericSignature=*/nullptr, extInfo, SILCoroutineKind::None,
-      ParameterConvention::Direct_Guaranteed, {getParamForArguments()},
+      ParameterConvention::Direct_Guaranteed,
+      {/*argumentBuffer=*/getRawPointerParmeter(),
+       /*resultBuffer=*/getRawPointerParmeter()},
       /*Yields=*/{},
-      /*Results=*/methodTy->getResults(),
+      /*Results=*/{},
       /*ErrorResult=*/methodTy->getErrorResult(),
       /*patternSubs=*/SubstitutionMap(),
       /*invocationSubs=*/SubstitutionMap(), Context);
@@ -257,30 +259,25 @@ void DistributedAccessor::computeArguments(llvm::Value *argumentBuffer,
 
 void DistributedAccessor::emit() {
   auto methodTy = Method->getLoweredFunctionType();
-  SILFunctionConventions conv(methodTy, IGF.getSILModule());
+  SILFunctionConventions targetConv(methodTy, IGF.getSILModule());
+  SILFunctionConventions accessorConv(AccessorType, IGF.getSILModule());
   TypeExpansionContext expansionContext = IGM.getMaximalTypeExpansionContext();
 
   auto params = IGF.collectParameters();
 
-  auto directResultTy = conv.getSILResultType(expansionContext);
+  auto directResultTy = targetConv.getSILResultType(expansionContext);
   const auto &directResultTI = IGM.getTypeInfo(directResultTy);
-  auto &resultSchema = directResultTI.nativeReturnValueSchema(IGM);
-  llvm::Value *indirectResultSlot = nullptr;
-
-  if (resultSchema.requiresIndirect())
-    indirectResultSlot = params.claimNext();
 
   Explosion arguments;
-  // Claim indirect results first, they are going to be passed
-  // through to the distributed method.
-  params.transferInto(arguments, conv.getNumIndirectSILResults());
 
   unsigned numAsyncContextParams =
       (unsigned)AsyncFunctionArgumentIndex::Context + 1;
   (void)params.claim(numAsyncContextParams);
 
   // UnsafeRawPointer that holds all of the argument values.
   auto *argBuffer = params.claimNext();
+  // UnsafeRawPointer that is used to store the result.
+  auto *resultBuffer = params.claimNext();
   // Reference to a `self` of the actor to be called.
   auto *actorSelf = params.claimNext();
 
@@ -298,6 +295,17 @@ void DistributedAccessor::emit() {
     emitAsyncFunctionPointer(IGM, IGF.CurFn, entity, AsyncLayout.getSize());
   }
 
+  auto *typedResultBuffer = IGF.Builder.CreateBitCast(
+    resultBuffer, IGM.getStoragePointerType(directResultTy));
+
+  if (targetConv.getNumIndirectSILResults()) {
+    // Since tuples are not allowed as valid result types (because they cannot
+    // conform to protocols), there could be only a single indirect result type
+    // associated with distributed method.
+    assert(targetConv.getNumIndirectSILResults() == 1);
+    arguments.add(typedResultBuffer);
+  }
+
   // Step one is to load all of the data from argument buffer,
   // so it could be forwarded to the distributed method.
   computeArguments(argBuffer, arguments);
@@ -316,30 +324,34 @@ void DistributedAccessor::emit() {
     emission->setArgs(arguments, /*isOutlined=*/false,
                       /*witnessMetadata=*/nullptr);
 
-    if (resultSchema.requiresIndirect()) {
-      Address resultAddr(indirectResultSlot,
+    // Load result of the thunk into the location provided by the caller.
+    // This would only generate code for direct results, if thunk has an
+    // indirect result (e.g. large struct) it result buffer would be passed
+    // as an argument.
+    {
+      Address resultAddr(typedResultBuffer,
                          directResultTI.getBestKnownAlignment());
       emission->emitToMemory(resultAddr, cast<LoadableTypeInfo>(directResultTI),
                              /*isOutlined=*/false);
-    } else {
-      emission->emitToExplosion(result, /*isOutlined=*/false);
     }
 
     // Both accessor and distributed method are always `async throws`
     // so we need to load error value (if any) from the slot.
     {
       assert(methodTy->hasErrorResult());
 
-      SILType errorType = conv.getSILErrorType(expansionContext);
+      SILType errorType = accessorConv.getSILErrorType(expansionContext);
       Address calleeErrorSlot =
           emission->getCalleeErrorSlot(errorType, /*isCalleeAsync=*/true);
       error.add(IGF.Builder.CreateLoad(calleeErrorSlot));
     }
 
     emission->end();
 
-    emitAsyncReturn(IGF, AsyncLayout, directResultTy, AccessorType, result,
-                    error);
+    Explosion voidResult;
+    emitAsyncReturn(IGF, AsyncLayout,
+                    accessorConv.getSILResultType(expansionContext),
+                    AccessorType, voidResult, error);
   }
 }
 
diff --git a/stdlib/public/Concurrency/Actor.cpp b/stdlib/public/Concurrency/Actor.cpp
@@ -2030,10 +2030,14 @@ SWIFT_CC(swiftasync)
 SWIFT_RUNTIME_STDLIB_SPI
 TargetExecutorSignature::FunctionType swift_distributed_execute_target;
 
-/// Accessor takes a context, an argument buffer as a raw pointer,
-/// and a reference to an actor.
-using DistributedAccessorSignature = AsyncSignature<void(void *, HeapObject *),
-                                                    /*throws=*/true>;
+/// Accessor takes:
+///   - an async context
+///   - an argument buffer as a raw pointer
+///   - a result buffer as a raw pointer
+///   - a reference to an actor to execute method on.
+using DistributedAccessorSignature =
+    AsyncSignature<void(void *, void *, HeapObject *),
+                   /*throws=*/true>;
 
 SWIFT_CC(swiftasync)
 static DistributedAccessorSignature::ContinuationType
@@ -2073,6 +2077,5 @@ void ::swift_distributed_execute_target(
   calleeContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
       swift_distributed_execute_target_resume);
 
-  // TODO: Add resultBuffer as an argument to store an indirect result into.
-  accessorEntry(calleeContext, argumentBuffer, actor);
+  accessorEntry(calleeContext, argumentBuffer, resultBuffer, actor);
 }
diff --git a/test/IRGen/distributed_actor_accessors.swift b/test/IRGen/distributed_actor_accessors.swift
@@ -139,7 +139,7 @@ public distributed actor MyOtherActor {
 
 // CHECK: define hidden swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiFTE"
 
-// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiFTETF"(%swift.context* swiftasync %0, i8* %1, %swift.refcounted* swiftself %2)
+// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiFTETF"(%swift.context* swiftasync %0, i8* %1, i8* %2, %swift.refcounted* swiftself %3)
 
 /// Read the current offset and cast an element to `Int`
 
@@ -192,14 +192,22 @@ public distributed actor MyOtherActor {
 // CHECK: [[STR_STRUCT:%.*]] = insertvalue { i64, %swift.bridge* } {{.*}}, %swift.bridge* {{.*}}, 1
 // CHECK: [[STR_SIZE:%.*]] = extractvalue { i64, %swift.bridge* } [[STR_STRUCT]], 0
 // CHECK-NEXT: [[STR_VAL:%.*]] = extractvalue { i64, %swift.bridge* } [[STR_STRUCT]], 1
-// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, i64 [[STR_SIZE]], %swift.bridge* [[STR_VAL]], %swift.error* {{.*}})
+
+/// Initialize the result buffer with values produced by the thunk
+
+// CHECK: store i64 [[STR_SIZE]], i64* %._guts._object._countAndFlagsBits._value, align 8
+// CHECK: store %swift.bridge* [[STR_VAL]], %swift.bridge** %._guts._object._object, align 8
+
+// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, %swift.error* {{.*}})
 
 /// ---> Thunk and distributed method accessor for `simple3`
 
 // CHECK: define hidden swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSFTE"
 
 /// !!! in `simple3` interesting bits are: argument value extraction (because string is exploded into N arguments) and call to distributed thunk
-// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSFTETF"
+// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSFTETF"(%swift.context* swiftasync {{.*}}, i8* [[ARG_BUFF:%.*]], i8* [[RESULT_BUFF:%.*]], %swift.refcounted* swiftself {{.*}})
+
+// CHECK: [[TYPED_RESULT_BUFF:%.*]] = bitcast i8* [[RESULT_BUFF]] to %TSi*
 
 // CHECK: %argval = load %TSS, %TSS* {{.*}}, align 8
 // CHECK: [[NATIVE_STR_PTR:%.*]] = bitcast %TSS* %argval.coercion.coerced to { i64, %swift.bridge* }*
@@ -220,13 +228,15 @@ public distributed actor MyOtherActor {
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i64, %swift.error* } [[THUNK_RESULT]], 0
 // CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
 // CHECK: [[INT_RES:%.*]] = extractvalue { i8*, i64, %swift.error* } [[THUNK_RESULT]], 1
-// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, i64 [[INT_RES]], %swift.error* {{.*}})
+// CHECK: %._value = getelementptr inbounds %TSi, %TSi* [[TYPED_RESULT_BUFF]], i32 0, i32 0
+// CHECK: store i64 [[INT_RES]], i64* %._value, align 8
+// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, %swift.error* {{.*}})
 
 /// --> Thunk and distributed method accessor for `single_case_enum`
 
 // CHECK: define hidden swifttailcc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFFTE"
 
-// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFFTETF"(%swift.context* swiftasync %0, i8* [[BUFFER:%.*]], %swift.refcounted* swiftself %2)
+// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFFTETF"(%swift.context* swiftasync %0, i8* [[BUFFER:%.*]], i8* [[RESULT_BUFF:%.*]], %swift.refcounted* swiftself {{.*}})
 
 /// First, let's check that there were no loads from the argument buffer and no stores to "current offset".
 
@@ -251,6 +261,7 @@ public distributed actor MyOtherActor {
 
 /// First, Load both arguments from the buffer.
 
+// CHECK: [[TYPED_RESULT_BUFF:%.*]] = bitcast i8* %2 to %T27distributed_actor_accessors9IndirectEO*
 // CHECK: store i8* %1, i8** %offset, align 8
 // CHECK-NEXT: %elt_offset = load i8*, i8** %offset, align 8
 // CHECK-NEXT: [[ENUM_PTR:%.*]] = bitcast i8* %elt_offset to %T27distributed_actor_accessors9IndirectEO*
@@ -261,7 +272,7 @@ public distributed actor MyOtherActor {
 // CHECK-NEXT: [[NATIVE_ENUM_VAL:%.*]] = load i64, i64* [[COERCED_ENUM_PTR]], align 8
 // CHECK: [[ENUM_PTR_INT:%.*]] = ptrtoint %T27distributed_actor_accessors9IndirectEO* [[ENUM_PTR]] to i64
 // CHECK-NEXT: [[NEXT_ELT_LOC:%.*]] = add i64 [[ENUM_PTR_INT]], 8
-// CHECK-NEXT: [[NEXT_ELT_PTR:%.*]] = inttoptr i64 %13 to i8*
+// CHECK-NEXT: [[NEXT_ELT_PTR:%.*]] = inttoptr i64 [[NEXT_ELT_LOC]] to i8*
 // CHECK-NEXT: store i8* [[NEXT_ELT_PTR]], i8** %offset, align 8
 // CHECK-NEXT: %elt_offset1 = load i8*, i8** %offset, align 8
 // CHECK-NEXT: [[INT_PTR:%.*]] = bitcast i8* %elt_offset1 to %TSi*
@@ -277,16 +288,22 @@ public distributed actor MyOtherActor {
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i64, %swift.error* } [[THUNK_RESULT]], 0
 // CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
 // CHECK: [[ENUM_RESULT:%.*]] = extractvalue { i8*, i64, %swift.error* } [[THUNK_RESULT]], 1
-// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, i64 [[ENUM_RESULT]], %swift.error* {{.*}})
+// CHECK: [[NATIVE_RESULT_PTR:%.*]] = bitcast %T27distributed_actor_accessors9IndirectEO* [[TYPED_RESULT_BUFF]] to i64*
+// CHECK-NEXT: store i64 [[ENUM_RESULT]], i64* [[NATIVE_RESULT_PTR]], align 8
+
+// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, %swift.error* {{.*}})
 
 /// ---> Thunk and distributed method for `complex`
 
 // CHECK: define hidden swifttailcc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtFTE"
 
-// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtFTETF"(%T27distributed_actor_accessors11LargeStructV* noalias nocapture [[INDIRECT_RES:%.*]], %swift.context* swiftasync %1, i8* %2, %swift.refcounted* swiftself %3)
+// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtFTETF"(%swift.context* swiftasync {{.*}}, i8* [[ARG_BUFF:%.*]], i8* [[RESULT_BUFF:%.*]], %swift.refcounted* swiftself {{.*}})
 
 /// First, let's check that all of the different argument types here are loaded correctly.
 
+/// Cast result buffer to the expected result type (in this case its indirect opaque pointer)
+// CHECK: [[TYPED_RESULT_BUFF:%.*]] = bitcast i8* [[RESULT_BUFF]] to %swift.opaque*
+
 /// -> [Int]
 
 // CHECK: %elt_offset = load i8*, i8** %offset, align 8
@@ -307,7 +324,7 @@ public distributed actor MyOtherActor {
 // CHECK-NEXT: [[NATIVE_OBJ_VAL:%.*]] = load %T27distributed_actor_accessors3ObjC*, %T27distributed_actor_accessors3ObjC** [[OBJ_PTR]], align 8
 // CHECK-NEXT: [[OBJ_PTR_INT:%.*]] = ptrtoint %T27distributed_actor_accessors3ObjC** [[OBJ_PTR]] to i64
 // CHECK-NEXT: [[NEXT_ELT:%.*]] = add i64 [[OBJ_PTR_INT]], 8
-// CHECK-NEXT: [[NEXT_ELT_PTR:%.*]] = inttoptr i64 %18 to i8*
+// CHECK-NEXT: [[NEXT_ELT_PTR:%.*]] = inttoptr i64 [[NEXT_ELT]] to i8*
 // CHECK-NEXT: store i8* [[NEXT_ELT_PTR]], i8** %offset, align 8
 
 /// -> String?
@@ -330,13 +347,16 @@ public distributed actor MyOtherActor {
 // CHECK-NEXT: %elt_offset5 = load i8*, i8** %offset, align 8
 // CHECK-NEXT: [[STRUCT_PTR:%.*]] = bitcast i8* %elt_offset5 to %T27distributed_actor_accessors11LargeStructV*
 // CHECK-NEXT: [[STRUCT_VAL:%.*]] = load %T27distributed_actor_accessors11LargeStructV, %T27distributed_actor_accessors11LargeStructV* [[STRUCT_PTR]], align 8
+
+// CHECK: [[INDIRECT_RESULT_BUFF:%.*]] = bitcast %swift.opaque* [[TYPED_RESULT_BUFF]] to %T27distributed_actor_accessors11LargeStructV*
+
 // CHECK: store %T27distributed_actor_accessors11LargeStructV [[STRUCT_VAL]], %T27distributed_actor_accessors11LargeStructV* %argval6.coercion.coerced, align 8
 // CHECK-NEXT: [[PTR_TO_STRUCT:%.*]] = bitcast %T27distributed_actor_accessors11LargeStructV* %argval6.coercion.coerced to %T27distributed_actor_accessors11LargeStructV**
 // CHECK-NEXT: [[NATIVE_STRUCT_VAL:%.*]] = load %T27distributed_actor_accessors11LargeStructV*, %T27distributed_actor_accessors11LargeStructV** [[PTR_TO_STRUCT]], align 8
 
 /// Now let's make sure that distributed thunk call uses the arguments correctly
 
-// CHECK: [[THUNK_RESULT:%.*]] = call { i8*, %swift.error* } (i32, i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8p0s_swift.errorss({{.*}}, %T27distributed_actor_accessors11LargeStructV* [[INDIRECT_RES]], %swift.context* {{.*}}, %swift.bridge* [[NATIVE_ARR_VAL]], %T27distributed_actor_accessors3ObjC* [[NATIVE_OBJ_VAL]], i64 [[NATIVE_OPT_VAL_0]], i64 [[NATIVE_OPT_VAL_1]], %T27distributed_actor_accessors11LargeStructV* [[NATIVE_STRUCT_VAL]], %T27distributed_actor_accessors7MyActorC* {{.*}})
+// CHECK: [[THUNK_RESULT:%.*]] = call { i8*, %swift.error* } (i32, i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8p0s_swift.errorss({{.*}}, %T27distributed_actor_accessors11LargeStructV* [[INDIRECT_RESULT_BUFF]], %swift.context* {{.*}}, %swift.bridge* [[NATIVE_ARR_VAL]], %T27distributed_actor_accessors3ObjC* [[NATIVE_OBJ_VAL]], i64 [[NATIVE_OPT_VAL_0]], i64 [[NATIVE_OPT_VAL_1]], %T27distributed_actor_accessors11LargeStructV* [[NATIVE_STRUCT_VAL]], %T27distributed_actor_accessors7MyActorC* {{.*}})
 
 /// RESULT is returned indirectly so there is nothing to pass to `end`
 
@@ -346,12 +366,13 @@ public distributed actor MyOtherActor {
 
 /// Let's check that there is no offset allocation here since parameter list is empty
 
-// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyFTETF"
+// CHECK: define internal swifttailcc void @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyFTETF"(%swift.context* swiftasync {{.*}}, i8* [[ARG_BUFF:%.*]], i8* [[RESULT_BUFF:%.*]], %swift.refcounted* swiftself {{.*}})
 // CHECK-NEXT: entry:
 // CHECK-NEXT: {{.*}} = alloca %swift.context*, align 8
 // CHECK-NEXT: %swifterror = alloca swifterror %swift.error*, align 8
 // CHECK-NEXT: {{.*}} = call token @llvm.coro.id.async(i32 20, i32 16, i32 0, i8* bitcast (%swift.async_func_pointer* @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyFTETFTu" to i8*))
-// CHECK-NEXT: {{.*}} = call i8* @llvm.coro.begin(token %4, i8* null)
+// CHECK-NEXT: {{.*}} = call i8* @llvm.coro.begin(token {{%.*}}, i8* null)
 // CHECK-NEXT: store %swift.context* {{.*}}, %swift.context** {{.*}}, align 8
 // CHECK-NEXT: store %swift.error* null, %swift.error** %swifterror, align 8
+// CHECK-NEXT: {{.*}} = bitcast i8* [[RESULT_BUFF]] to %swift.opaque*
 // CHECK-NEXT: {{.*}} = load i32, i32* getelementptr inbounds (%swift.async_func_pointer, %swift.async_func_pointer* bitcast (void (%swift.context*, %T27distributed_actor_accessors12MyOtherActorC*)* @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyFTE" to %swift.async_func_pointer*), i32 0, i32 0), align 8
