Merge pull request #77028 from gottesmm/swift_task_runOnMainActor

[concurrency] Implement swift_task_runOnMainActor/hop to main actor thunk

Author: gottesmm

include/swift/AST/ASTMangler.h

@@ -15,6 +15,7 @@
 
 #include "swift/AST/Decl.h"
 #include "swift/AST/FreestandingMacroExpansion.h"
+#include "swift/AST/SILThunkKind.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/Mangler.h"
 #include "swift/Basic/TaggedUnion.h"
@@ -314,6 +315,9 @@ class ASTMangler : public Mangler {
                                                 DifferentiabilityKind kind,
                                                 const AutoDiffConfig &config);
 
+  std::string mangleSILThunkKind(StringRef originalName,
+                                 SILThunkKind thunkKind);
+
   /// Mangle the AutoDiff generated declaration for the given:
   /// - Generated declaration kind: linear map struct or branching trace enum.
   /// - Mangled original function name.

include/swift/AST/SILThunkKind.h

@@ -0,0 +1,84 @@
+//===--- SILThunkKind.h ---------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file This file contains a kind that defines the types of thunks that can be
+/// generated using a thunk inst. It provides an AST level interface that lets
+/// one generate the derived function kind and is also used to make adding such
+/// kinds to the mangler trivial.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_SILTHUNKKIND_H
+#define SWIFT_AST_SILTHUNKKIND_H
+
+#include "swift/AST/Types.h"
+
+namespace swift {
+
+class SILType;
+
+struct SILThunkKind {
+  enum InnerTy {
+    Invalid = 0,
+
+    /// A thunk that just calls the passed in function. Used for testing
+    /// purposes of the underlying thunking machinery.
+    Identity = 1,
+
+    /// A thunk that checks if a value is on the main actor and if it isn't
+    /// jumps to the main actor. It expects that the passed in function does not
+    /// have any arguments, results, is synchronous, and does not throw.
+    HopToMainActorIfNeeded = 2,
+
+    MaxValue = HopToMainActorIfNeeded,
+  };
+
+  InnerTy innerTy;
+
+  SILThunkKind() : innerTy(InnerTy::Invalid) {}
+  SILThunkKind(InnerTy innerTy) : innerTy(innerTy) {}
+  SILThunkKind(unsigned inputInnerTy) : innerTy(InnerTy(inputInnerTy)) {
+    assert(inputInnerTy <= MaxValue && "Invalid value");
+  }
+
+  operator InnerTy() const { return innerTy; }
+
+  /// Given the current enum state returned the derived output function from
+  /// \p inputFunction.
+  ///
+  /// Defined in Instructions.cpp
+  CanSILFunctionType getDerivedFunctionType(SILFunction *fn,
+                                            CanSILFunctionType inputFunction,
+                                            SubstitutionMap subMap) const;
+
+  /// Given the current enum state returned the derived output function from
+  /// \p inputFunction.
+  ///
+  /// Defined in Instructions.cpp
+  SILType getDerivedFunctionType(SILFunction *fn, SILType inputFunctionType,
+                                 SubstitutionMap subMap) const;
+
+  Demangle::MangledSILThunkKind getMangledKind() const {
+    switch (innerTy) {
+    case Invalid:
+      return Demangle::MangledSILThunkKind::Invalid;
+    case Identity:
+      return Demangle::MangledSILThunkKind::Identity;
+    case HopToMainActorIfNeeded:
+      return Demangle::MangledSILThunkKind::HopToMainActorIfNeeded;
+    }
+  }
+};
+
+} // namespace swift
+
+#endif

include/swift/Basic/Features.def

@@ -433,6 +433,9 @@ EXPERIMENTAL_FEATURE(CoroutineAccessorsUnwindOnCallerError, false)
 /// modify/read coroutines use the callee-allocated ABI
 EXPERIMENTAL_FEATURE(CoroutineAccessorsAllocateInCallee, false)
 
+// When a parameter has unspecified isolation, infer it as main actor isolated.
+EXPERIMENTAL_FEATURE(GenerateForceToMainActorThunks, false)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE

include/swift/Demangling/Demangle.h

@@ -136,6 +136,12 @@ enum class MangledDifferentiabilityKind : char {
   Linear = 'l',
 };
 
+enum class MangledSILThunkKind : char {
+  Invalid = 0,
+  Identity = 'I',
+  HopToMainActorIfNeeded = 'H',
+};
+
 /// The pass that caused the specialization to occur. We use this to make sure
 /// that two passes that generate similar changes do not yield the same
 /// mangling. This currently cannot happen, so this is just a safety measure

include/swift/Demangling/DemangleNodes.def

@@ -291,6 +291,8 @@ NODE(ReflectionMetadataAssocTypeDescriptor)
 NODE(ReflectionMetadataSuperclassDescriptor)
 NODE(GenericTypeParamDecl)
 NODE(CurryThunk)
+NODE(SILThunkIdentity)
+NODE(SILThunkHopToMainActorIfNeeded)
 NODE(DispatchThunk)
 NODE(MethodDescriptor)
 NODE(ProtocolRequirementsBaseDescriptor)

include/swift/Runtime/Concurrency.h

@@ -963,6 +963,45 @@ TaskExecutorRef swift_task_getPreferredTaskExecutor(void);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_task_isCurrentExecutor(SerialExecutorRef executor);
 
+/// This is an options enum that is used to pass flags to
+/// swift_task_isCurrentExecutorWithFlags. It is meant to be a flexible toggle.
+///
+/// Since this is an options enum, so all values should be powers of 2.
+///
+/// NOTE: We are purposely leaving this as a uint64_t so that on all platforms
+/// this could be a pointer to a different enum instance if we need it to be.
+enum swift_task_is_current_executor_flag : uint64_t {
+  /// We aren't passing any flags.
+  None = 0x0,
+
+  /// This is not used today, but is just future ABI reservation.
+  ///
+  /// The intention is that we may want the ability to tell future versions of
+  /// the runtime that this uint64_t is actually a pointer that it should
+  /// dereference and then have further extended behavior controlled by a
+  /// different enum. By placing this here, we ensure that we will have a tagged
+  /// pointer compatible flag for this purpose.
+  TaggedPointer = 0x1,
+
+  /// This is not used today, but is just future ABI reservation.
+  ///
+  /// \see swift_task_is_current_executor_flag::TaggedPointer
+  TaggedPointer2 = 0x2,
+
+  /// This is not used today, but is just future ABI reservation.
+  ///
+  /// \see swift_task_is_current_executor_flag::TaggedPointer
+  TaggedPointer3 = 0x4,
+
+  /// The routine should assert on failure.
+  Assert = 0x8,
+};
+
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift)
+bool swift_task_isCurrentExecutorWithFlags(
+    SerialExecutorRef executor, swift_task_is_current_executor_flag flags);
+
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_reportUnexpectedExecutor(
     const unsigned char *file, uintptr_t fileLength, bool fileIsASCII,

include/swift/SIL/SILBuilder.h

@@ -1251,6 +1251,13 @@ class SILBuilder {
     return createThunk(Loc, Op, ThunkInst::Kind::Identity, substitutionMap);
   }
 
+  ThunkInst *
+  createHopToMainActorIfNeededThunk(SILLocation Loc, SILValue Op,
+                                    SubstitutionMap substitutionMap = {}) {
+    return createThunk(Loc, Op, ThunkInst::Kind::HopToMainActorIfNeeded,
+                       substitutionMap);
+  }
+
   ConvertEscapeToNoEscapeInst *
   createConvertEscapeToNoEscape(SILLocation Loc, SILValue Op, SILType Ty,
                                 bool lifetimeGuaranteed) {

include/swift/SIL/SILInstruction.h

@@ -23,6 +23,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/ProtocolConformanceRef.h"
+#include "swift/AST/SILThunkKind.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/TypeAlignments.h"
 #include "swift/Basic/Compiler.h"
@@ -5887,41 +5888,7 @@ class ThunkInst final
     : public UnaryInstructionWithTypeDependentOperandsBase<
           SILInstructionKind::ThunkInst, ThunkInst, SingleValueInstruction> {
 public:
-  struct Kind {
-    enum InnerTy {
-      Invalid = 0,
-
-      /// A thunk that just calls the passed in function. Used for testing
-      /// purposes of the underlying thunking machinery.
-      Identity = 1,
-
-      MaxValue = Identity,
-    };
-
-    InnerTy innerTy;
-
-    Kind() : innerTy(InnerTy::Invalid) {}
-    Kind(InnerTy innerTy) : innerTy(innerTy) {}
-    Kind(unsigned inputInnerTy) : innerTy(InnerTy(inputInnerTy)) {
-      assert(inputInnerTy <= MaxValue && "Invalid value");
-    }
-
-    operator InnerTy() const { return innerTy; }
-
-    /// Given the current enum state returned the derived output function from
-    /// \p inputFunction.
-    CanSILFunctionType getDerivedFunctionType(SILFunction *fn,
-                                              CanSILFunctionType inputFunction,
-                                              SubstitutionMap subMap) const;
-
-    SILType getDerivedFunctionType(SILFunction *fn, SILType inputFunctionType,
-                                   SubstitutionMap subMap) const {
-      auto fType = inputFunctionType.castTo<SILFunctionType>();
-      return SILType::getPrimitiveType(
-          getDerivedFunctionType(fn, fType, subMap),
-          inputFunctionType.getCategory());
-    }
-  };
+  using Kind = SILThunkKind;
 
   /// The type of thunk we are supposed to produce.
   Kind kind;

lib/AST/ASTMangler.cpp

@@ -636,6 +636,16 @@ std::string ASTMangler::mangleSILDifferentiabilityWitness(StringRef originalName
   return finalize();
 }
 
+std::string ASTMangler::mangleSILThunkKind(StringRef originalName,
+                                           SILThunkKind thunkKind) {
+  beginManglingWithoutPrefix();
+  appendOperator(originalName);
+  // Prefix for thunk inst based thunks
+  auto code = (char)thunkKind.getMangledKind();
+  appendOperator("TT", StringRef(&code, 1));
+  return finalize();
+}
+
 std::string ASTMangler::mangleAutoDiffGeneratedDeclaration(
     AutoDiffGeneratedDeclarationKind declKind, StringRef origFnName,
     unsigned bbId, AutoDiffLinearMapKind linearMapKind,

lib/AST/FeatureSet.cpp

@@ -199,6 +199,7 @@ UNINTERESTING_FEATURE(SameElementRequirements)
 UNINTERESTING_FEATURE(UnspecifiedMeansMainActorIsolated)
 UNINTERESTING_FEATURE(GlobalActorInferenceCutoff)
 UNINTERESTING_FEATURE(KeyPathWithStaticMembers)
+UNINTERESTING_FEATURE(GenerateForceToMainActorThunks)
 
 static bool usesFeatureSendingArgsAndResults(Decl *decl) {
   auto isFunctionTypeWithSending = [](Type type) {