[Distributed] IRGen: Teach distributed thunk accessor to load witness tables (when needed)

Author: xedin

lib/IRGen/GenDistributed.cpp

@@ -121,6 +121,16 @@ class DistributedAccessor {
                const SILParameterInfo &param,
                Explosion &arguments);
 
+  /// Load witness table addresses (if any) from the given buffer
+  /// into the given argument explosion.
+  ///
+  /// Number of witnesses to load is provided by \c numTables but
+  /// it's checked against the number of \c expectedWitnessTables.
+  void emitLoadOfWitnessTables(llvm::Value *witnessTables,
+                               llvm::Value *numTables,
+                               unsigned expectedWitnessTables,
+                               Explosion &arguments);
+
   FunctionPointer getPointerToTarget() const;
 
   Callee getCalleeForDistributedTarget(llvm::Value *self) const;
@@ -140,6 +150,11 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
                             ParameterConvention::Direct_Guaranteed);
   };
 
+  auto getUIntParameter = [&]() {
+    return SILParameterInfo(Context.getUIntType()->getCanonicalType(),
+                            ParameterConvention::Direct_Guaranteed);
+  };
+
   // `self` of the distributed actor is going to be passed as an argument
   // to this accessor function.
   auto extInfo = SILExtInfoBuilder()
@@ -162,6 +177,8 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
        /*argumentTypes=*/getRawPointerParameter(),
        /*resultBuffer=*/getRawPointerParameter(),
        /*substitutions=*/getRawPointerParameter(),
+       /*witnessTables=*/getRawPointerParameter(),
+       /*numWitnessTables=*/getUIntParameter(),
        /*actor=*/targetTy->getParameters().back()},
       /*Yields=*/{},
       /*Results=*/{},
@@ -432,6 +449,37 @@ DistributedAccessor::loadArgument(unsigned argumentIdx,
   return {alignedOffset.getAddress(), typeInfo.getSize(IGF, paramTy)};
 }
 
+void DistributedAccessor::emitLoadOfWitnessTables(llvm::Value *witnessTables,
+                                                  llvm::Value *numTables,
+                                                  unsigned expectedWitnessTables,
+                                                  Explosion &arguments) {
+  auto contBB = IGF.createBasicBlock("");
+  auto unreachableBB = IGF.createBasicBlock("incorrect-witness-tables");
+
+  auto incorrectNum = IGF.Builder.CreateICmpNE(
+      numTables, llvm::ConstantInt::get(IGM.SizeTy, expectedWitnessTables));
+
+  // Make sure that we have a correct number of witness tables provided to us.
+  IGF.Builder.CreateCondBr(incorrectNum, unreachableBB, contBB);
+  {
+    IGF.Builder.emitBlock(unreachableBB);
+    IGF.Builder.CreateUnreachable();
+  }
+
+  IGF.Builder.emitBlock(contBB);
+
+  witnessTables = IGF.Builder.CreateBitCast(witnessTables, IGM.Int8PtrPtrTy);
+
+  for (unsigned i = 0, n = expectedWitnessTables; i != n; ++i) {
+    auto offset = Size(i * IGM.getPointerSize());
+    auto alignment = IGM.getPointerAlignment();
+
+    auto witnessTableAddr = IGF.emitAddressAtOffset(
+        witnessTables, Offset(offset), IGM.Int8PtrTy, Alignment(alignment));
+    arguments.add(witnessTableAddr.getAddress());
+  }
+}
+
 void DistributedAccessor::emit() {
   auto targetTy = Target->getLoweredFunctionType();
   SILFunctionConventions targetConv(targetTy, IGF.getSILModule());
@@ -457,6 +505,10 @@ void DistributedAccessor::emit() {
   auto *resultBuffer = params.claimNext();
   // UnsafeRawPointer that represents a list of substitutions
   auto *substitutions = params.claimNext();
+  // UnsafeRawPointer that represents a list of witness tables
+  auto *witnessTables = params.claimNext();
+  // Integer that represented the number of witness tables
+  auto *numWitnessTables = params.claimNext();
   // Reference to a `self` of the actor to be called.
   auto *actorSelf = params.claimNext();
 
@@ -490,12 +542,22 @@ void DistributedAccessor::emit() {
   computeArguments(argBuffer, argTypes, arguments);
 
   // Add all of the substitutions to the explosion
-  if (auto *environment = Target->getGenericEnvironment()) {
+  if (auto *genericEnvironment = Target->getGenericEnvironment()) {
     // swift.type **
     llvm::Value *substitutionBuffer =
         IGF.Builder.CreateBitCast(substitutions, IGM.TypeMetadataPtrPtrTy);
 
-    for (unsigned index : indices(environment->getGenericParams())) {
+    // Collect the generic arguments expected by the distributed thunk.
+    // We need this to determine the expected number of witness tables
+    // to load from the buffer provided by the caller.
+    llvm::SmallVector<llvm::Type *, 4> targetGenericArguments;
+    expandPolymorphicSignature(IGM, targetTy, targetGenericArguments);
+
+    unsigned numGenericArgs = genericEnvironment->getGenericParams().size();
+    unsigned expectedWitnessTables =
+        targetGenericArguments.size() - numGenericArgs;
+
+    for (unsigned index = 0; index < numGenericArgs; ++index) {
       auto offset =
           Size(index * IGM.DataLayout.getTypeAllocSize(IGM.TypeMetadataPtrTy));
       auto alignment =
@@ -506,6 +568,9 @@ void DistributedAccessor::emit() {
                                   IGM.TypeMetadataPtrTy, Alignment(alignment));
       arguments.add(IGF.Builder.CreateLoad(substitution, "substitution"));
     }
+
+    emitLoadOfWitnessTables(witnessTables, numWitnessTables,
+                            expectedWitnessTables, arguments);
   }
 
   // Step two, let's form and emit a call to the distributed method

stdlib/public/Concurrency/Actor.cpp

@@ -1724,7 +1724,9 @@ void *swift_distributed_get_generic_environment(const char *targetNameStart,
 ///    argumentBuffer: Builtin.RawPointer,
 ///    argumentTypes: UnsafeBufferPointer<Any.Type>,
 ///    resultBuffer: Builtin.RawPointer,
-///    substitutions: UnsafeRawPointer?
+///    substitutions: UnsafeRawPointer?,
+///    witnessTables: UnsafeRawPointer?,
+///    numWitnessTables: UInt
 /// ) async throws
 using TargetExecutorSignature =
     AsyncSignature<void(/*on=*/DefaultActor *,
@@ -1733,6 +1735,8 @@ using TargetExecutorSignature =
                         /*argumentTypes=*/const Metadata *const *,
                         /*resultBuffer=*/void *,
                         /*substitutions=*/void *,
+                        /*witnessTables=*/void **,
+                        /*numWitnessTables=*/size_t,
                         /*resumeFunc=*/TaskContinuationFunction *,
                         /*callContext=*/AsyncContext *),
                    /*throws=*/true>;
@@ -1747,12 +1751,16 @@ TargetExecutorSignature::FunctionType swift_distributed_execute_target;
 ///   - a list of all argument types (with substitutions applied)
 ///   - a result buffer as a raw pointer
 ///   - a list of substitutions
+///   - a list of witness tables
+///   - a number of witness tables in the buffer
 ///   - a reference to an actor to execute method on.
 using DistributedAccessorSignature =
     AsyncSignature<void(/*argumentBuffer=*/void *,
                         /*argumentTypes=*/const Metadata *const *,
                         /*resultBuffer=*/void *,
                         /*substitutions=*/void *,
+                        /*witnessTables=*/void **,
+                        /*numWitnessTables=*/size_t,
                         /*actor=*/HeapObject *),
                    /*throws=*/true>;
 
@@ -1783,6 +1791,8 @@ void ::swift_distributed_execute_target(
     const Metadata *const *argumentTypes,
     void *resultBuffer,
     void *substitutions,
+    void **witnessTables,
+    size_t numWitnessTables,
     TaskContinuationFunction *resumeFunc,
     AsyncContext *callContext) {
   auto *accessor = findDistributedAccessor(targetNameStart, targetNameLength);
@@ -1820,5 +1830,7 @@ void ::swift_distributed_execute_target(
                 argumentBuffer, argumentTypes,
                 resultBuffer,
                 substitutions,
+                witnessTables,
+                numWitnessTables,
                 actor);
 }

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -324,7 +324,9 @@ extension DistributedActorSystem {
         argumentBuffer: hargs.buffer._rawValue, // TODO(distributed): pass the invocationDecoder instead, so we can decode inside IRGen directly into the argument explosion
         argumentTypes: argumentTypesBuffer.baseAddress!._rawValue,
         resultBuffer: resultBuffer._rawValue,
-        substitutions: UnsafeRawPointer(substitutionsBuffer)
+        substitutions: UnsafeRawPointer(substitutionsBuffer),
+        witnessTables: witnessTablesBuffer,
+        numWitnessTables: UInt(numWitnessTables)
       )
 
       func onReturn<R>(_ resultTy: R.Type) async throws {
@@ -345,7 +347,9 @@ func _executeDistributedTarget(
   argumentBuffer: Builtin.RawPointer, // HeterogeneousBuffer of arguments
   argumentTypes: Builtin.RawPointer,
   resultBuffer: Builtin.RawPointer,
-  substitutions: UnsafeRawPointer?
+  substitutions: UnsafeRawPointer?,
+  witnessTables: UnsafeRawPointer?,
+  numWitnessTables: UInt
 ) async throws
 
 // ==== ----------------------------------------------------------------------------------------------------------------

test/Distributed/Runtime/distributed_actor_remoteCall.swift

@@ -34,10 +34,19 @@ enum E : Sendable, Codable {
   case foo, bar
 }
 
+struct S<T: Codable> : Codable {
+  var data: T
+}
+
 @_silgen_name("swift_distributed_actor_is_remote")
 func __isRemoteActor(_ actor: AnyObject) -> Bool
 
 distributed actor Greeter {
+  distributed func generic1<T: Codable, U: Codable>(t: T, u: U) {
+    print("---> T = \(t), type(of:) = \(type(of: t))")
+    print("---> U = \(u), type(of:) = \(type(of: u))")
+  }
+
   distributed func empty() {
   }
 
@@ -60,7 +69,6 @@ distributed actor Greeter {
   distributed func enumResult() -> E {
     .bar
   }
-
 }
 
 
@@ -158,7 +166,6 @@ struct FakeInvocation: DistributedTargetInvocationEncoder, DistributedTargetInvo
       fatalError("Cannot cast argument\(anyArgument) to expected \(Argument.self)")
     }
 
-    print("  > decode argument: \(argument)")
     pointer.initialize(to: argument)
     argumentIndex += 1
   }
@@ -194,6 +201,8 @@ let answerName = "$s4main7GreeterC6answerSiyFTE"
 let largeResultName = "$s4main7GreeterC11largeResultAA11LargeStructVyFTE"
 let enumResultName = "$s4main7GreeterC10enumResultAA1EOyFTE"
 let echoName = "$s4main7GreeterC4echo4name3ageS2S_SitFTE"
+// <T: Codable, U: Codable>(t: T, u: U)
+let generic1Name = "$s4main7GreeterC8generic11t1uyx_q_tSeRzSERzSeR_SER_r0_lFTE"
 
 func test() async throws {
   let system = FakeActorSystem()
@@ -255,6 +264,25 @@ func test() async throws {
   )
   // CHECK: RETURN: Echo: name: Caplin, age: 42
 
+  var generic1Invocation = system.makeInvocationEncoder()
+
+  try generic1Invocation.recordGenericSubstitution(Int.self)
+  try generic1Invocation.recordGenericSubstitution(String.self)
+  try generic1Invocation.recordArgument(42)
+  try generic1Invocation.recordArgument("Ultimate Question!")
+  try generic1Invocation.doneRecording()
+
+  try await system.executeDistributedTarget(
+    on: local,
+    mangledTargetName: generic1Name,
+    invocationDecoder: &generic1Invocation,
+    handler: FakeResultHandler()
+  )
+
+  // CHECK: ---> T = 42, type(of:) = Int
+  // CHECK-NEXT: ---> U = Ultimate Question!, type(of:) = String
+  // CHECK-NEXT: RETURN: ()
+
   print("done")
   // CHECK-NEXT: done
 }