Fix function differentiability kind metadata and mangling. (#36601)

* Move differentiability kinds from target function type metadata to trailing objects so that we don't exhaust all remaining bits of function type metadata.
  * Differentiability kind is now stored in a tail-allocated word when function type flags say it's differentiable, located immediately after the normal function type metadata's contents (with proper alignment in between).
  * Add new runtime function `swift_getFunctionTypeMetadataDifferentiable` which handles differentiable function types.
* Fix mangling of different differentiability kinds in function types. Mangle it like `ConcurrentFunctionType` so that we can drop special cases for escaping functions.
    ```
    function-signature ::= params-type params-type async? sendable? throws? differentiable? // results and parameters
    ...
    differentiable ::= 'jf'                    // @differentiable(_forward) on function type
    differentiable ::= 'jr'                    // @differentiable(reverse) on function type
    differentiable ::= 'jd'                    // @differentiable on function type
    differentiable ::= 'jl'                    // @differentiable(_linear) on function type
    ```

Resolves rdar://75240064.

Author: rxwei

docs/ABI/Mangling.rst

@@ -560,24 +560,24 @@ Types
   FUNCTION-KIND ::= 'zC' C-TYPE              // C function pointer type with with non-canonical C type
   FUNCTION-KIND ::= 'A'                      // @auto_closure function type (escaping)
   FUNCTION-KIND ::= 'E'                      // function type (noescape)
-  FUNCTION-KIND ::= 'F'                      // @differentiable function type
-  FUNCTION-KIND ::= 'G'                      // @differentiable function type (escaping)
-  FUNCTION-KIND ::= 'H'                      // @differentiable(_linear) function type
-  FUNCTION-KIND ::= 'I'                      // @differentiable(_linear) function type (escaping)
 
   C-TYPE is mangled according to the Itanium ABI, and prefixed with the length.
   Non-ASCII identifiers are preserved as-is; we do not use Punycode.
 
-  function-signature ::= params-type params-type async? sendable? throws? // results and parameters
+  function-signature ::= params-type params-type async? sendable? throws? differentiable? // results and parameters
 
-  params-type ::= type 'z'? 'h'?              // tuple in case of multiple parameters or a single parameter with a single tuple type
+  params-type ::= type 'z'? 'h'?             // tuple in case of multiple parameters or a single parameter with a single tuple type
                                              // with optional inout convention, shared convention. parameters don't have labels,
                                              // they are mangled separately as part of the entity.
-  params-type ::= empty-list                  // shortcut for no parameters
+  params-type ::= empty-list                 // shortcut for no parameters
 
-  sendable ::= 'J'                         // @Sendable on function types
+  sendable ::= 'J'                           // @Sendable on function types
   async ::= 'Y'                              // 'async' annotation on function types
   throws ::= 'K'                             // 'throws' annotation on function types
+  differentiable ::= 'jf'                    // @differentiable(_forward) on function type
+  differentiable ::= 'jr'                    // @differentiable(reverse) on function type
+  differentiable ::= 'jd'                    // @differentiable on function type
+  differentiable ::= 'jl'                    // @differentiable(_linear) on function type
 
   type-list ::= list-type '_' list-type*     // list of types
   type-list ::= empty-list

include/swift/ABI/Metadata.h

@@ -1658,6 +1658,7 @@ struct TargetFunctionTypeMetadata : public TargetMetadata<Runtime> {
   bool isAsync() const { return Flags.isAsync(); }
   bool isThrowing() const { return Flags.isThrowing(); }
   bool isSendable() const { return Flags.isSendable(); }
+  bool isDifferentiable() const { return Flags.isDifferentiable(); }
   bool hasParameterFlags() const { return Flags.hasParameterFlags(); }
   bool isEscaping() const { return Flags.isEscaping(); }
 
@@ -1675,6 +1676,28 @@ struct TargetFunctionTypeMetadata : public TargetMetadata<Runtime> {
     return reinterpret_cast<const uint32_t *>(getParameters() +
                                               getNumParameters());
   }
+
+  TargetFunctionMetadataDifferentiabilityKind<StoredSize> *
+  getDifferentiabilityKindAddress() {
+    assert(isDifferentiable());
+    void *previousEndAddr = hasParameterFlags()
+        ? reinterpret_cast<void *>(getParameterFlags() + getNumParameters())
+        : reinterpret_cast<void *>(getParameters() + getNumParameters());
+    return reinterpret_cast<
+        TargetFunctionMetadataDifferentiabilityKind<StoredSize> *>(
+        llvm::alignAddr(previousEndAddr,
+                        llvm::Align(alignof(typename Runtime::StoredPointer))));
+  }
+
+  TargetFunctionMetadataDifferentiabilityKind<StoredSize>
+  getDifferentiabilityKind() const {
+    if (isDifferentiable()) {
+      return *const_cast<TargetFunctionTypeMetadata<Runtime> *>(this)
+          ->getDifferentiabilityKindAddress();
+    }
+    return TargetFunctionMetadataDifferentiabilityKind<StoredSize>
+        ::NonDifferentiable;
+  }
 };
 using FunctionTypeMetadata = TargetFunctionTypeMetadata<InProcess>;
 

include/swift/ABI/MetadataValues.h

@@ -787,13 +787,29 @@ enum class FunctionMetadataConvention: uint8_t {
 
 /// Differentiability kind for function type metadata.
 /// Duplicates `DifferentiabilityKind` in AST/AutoDiff.h.
-enum class FunctionMetadataDifferentiabilityKind: uint8_t {
-  NonDifferentiable = 0b00000,
-  Forward           = 0b00001,
-  Reverse           = 0b00010,
-  Normal            = 0b00011,
-  Linear            = 0b10000,
+template <typename int_type>
+struct TargetFunctionMetadataDifferentiabilityKind {
+  enum Value : int_type {
+    NonDifferentiable = 0,
+    Forward           = 1,
+    Reverse           = 2,
+    Normal            = 3,
+    Linear            = 4,
+  } Value;
+
+  constexpr TargetFunctionMetadataDifferentiabilityKind(
+      enum Value value = NonDifferentiable) : Value(value) {}
+
+  int_type getIntValue() const {
+    return (int_type)Value;
+  }
+
+  bool isDifferentiable() const {
+    return Value != NonDifferentiable;
+  }
 };
+using FunctionMetadataDifferentiabilityKind =
+    TargetFunctionMetadataDifferentiabilityKind<size_t>;
 
 /// Flags in a function type metadata record.
 template <typename int_type>
@@ -808,8 +824,7 @@ class TargetFunctionTypeFlags {
     ThrowsMask             = 0x01000000U,
     ParamFlagsMask         = 0x02000000U,
     EscapingMask           = 0x04000000U,
-    DifferentiabilityMask  = 0x98000000U,
-    DifferentiabilityShift = 27U,
+    DifferentiableMask     = 0x08000000U,
     AsyncMask              = 0x20000000U,
     SendableMask           = 0x40000000U,
   };
@@ -842,10 +857,10 @@ class TargetFunctionTypeFlags {
                                              (throws ? ThrowsMask : 0));
   }
 
-  constexpr TargetFunctionTypeFlags<int_type> withDifferentiabilityKind(
-      FunctionMetadataDifferentiabilityKind differentiabilityKind) const {
-    return TargetFunctionTypeFlags((Data & ~DifferentiabilityMask)
-        | (int_type(differentiabilityKind) << DifferentiabilityShift));
+  constexpr TargetFunctionTypeFlags<int_type>
+  withDifferentiable(bool differentiable) const {
+    return TargetFunctionTypeFlags<int_type>((Data & ~DifferentiableMask) |
+                                     (differentiable ? DifferentiableMask : 0));
   }
 
   constexpr TargetFunctionTypeFlags<int_type>
@@ -888,13 +903,7 @@ class TargetFunctionTypeFlags {
   bool hasParameterFlags() const { return bool(Data & ParamFlagsMask); }
 
   bool isDifferentiable() const {
-    return getDifferentiabilityKind() !=
-        FunctionMetadataDifferentiabilityKind::NonDifferentiable;
-  }
-
-  FunctionMetadataDifferentiabilityKind getDifferentiabilityKind() const {
-    return FunctionMetadataDifferentiabilityKind(
-        (Data & DifferentiabilityMask) >> DifferentiabilityShift);
+    return bool (Data & DifferentiableMask);
   }
 
   int_type getIntValue() const {
@@ -947,6 +956,12 @@ class TargetParameterTypeFlags {
         (Data & ~AutoClosureMask) | (isAutoClosure ? AutoClosureMask : 0));
   }
 
+  constexpr TargetParameterTypeFlags<int_type>
+  withNoDerivative(bool isNoDerivative) const {
+    return TargetParameterTypeFlags<int_type>(
+        (Data & ~NoDerivativeMask) | (isNoDerivative ? NoDerivativeMask : 0));
+  }
+
   bool isNone() const { return Data == 0; }
   bool isVariadic() const { return Data & VariadicMask; }
   bool isAutoClosure() const { return Data & AutoClosureMask; }

include/swift/AST/ASTDemangler.h

@@ -97,8 +97,10 @@ class ASTBuilder {
 
   Type createTupleType(ArrayRef<Type> eltTypes, StringRef labels);
 
-  Type createFunctionType(ArrayRef<Demangle::FunctionParam<Type>> params,
-                          Type output, FunctionTypeFlags flags);
+  Type createFunctionType(
+      ArrayRef<Demangle::FunctionParam<Type>> params,
+      Type output, FunctionTypeFlags flags,
+      FunctionMetadataDifferentiabilityKind diffKind);
 
   Type createImplFunctionType(
       Demangle::ImplParameterConvention calleeConvention,

include/swift/Demangling/DemangleNodes.def

@@ -69,10 +69,6 @@ NODE(DependentProtocolConformanceInherited)
 NODE(DependentProtocolConformanceAssociated)
 CONTEXT_NODE(Destructor)
 CONTEXT_NODE(DidSet)
-NODE(DifferentiableFunctionType)
-NODE(EscapingDifferentiableFunctionType)
-NODE(LinearFunctionType)
-NODE(EscapingLinearFunctionType)
 NODE(Directness)
 NODE(DynamicAttribute)
 NODE(DirectMethodReferenceAttribute)
@@ -86,6 +82,7 @@ NODE(ErrorType)
 NODE(EscapingAutoClosureType)
 NODE(NoEscapeFunctionType)
 NODE(ConcurrentFunctionType)
+NODE(DifferentiableFunctionType)
 NODE(ExistentialMetatype)
 CONTEXT_NODE(ExplicitClosure)
 CONTEXT_NODE(Extension)

include/swift/Demangling/Demangler.h

@@ -575,6 +575,7 @@ class Demangler : public NodeFactory {
   NodePointer demangleAutoDiffSelfReorderingReabstractionThunk();
   NodePointer demangleDifferentiabilityWitness();
   NodePointer demangleIndexSubset();
+  NodePointer demangleDifferentiableFunctionType();
 
   bool demangleBoundGenerics(Vector<NodePointer> &TypeListList,
                              NodePointer &RetroactiveConformances);

include/swift/Demangling/TypeDecoder.h

@@ -316,6 +316,11 @@ class ImplFunctionTypeFlags {
 
   bool isPseudogeneric() const { return Pseudogeneric; }
 
+  bool isDifferentiable() const {
+    return getDifferentiabilityKind() !=
+        ImplFunctionDifferentiabilityKind::NonDifferentiable;
+  }
+
   ImplFunctionDifferentiabilityKind getDifferentiabilityKind() const {
     return ImplFunctionDifferentiabilityKind(DifferentiabilityKind);
   }
@@ -708,10 +713,6 @@ class TypeDecoder {
     case NodeKind::NoEscapeFunctionType:
     case NodeKind::AutoClosureType:
     case NodeKind::EscapingAutoClosureType:
-    case NodeKind::DifferentiableFunctionType:
-    case NodeKind::EscapingDifferentiableFunctionType:
-    case NodeKind::LinearFunctionType:
-    case NodeKind::EscapingLinearFunctionType:
     case NodeKind::FunctionType: {
       if (Node->getNumChildren() < 2)
         return MAKE_NODE_TYPE_ERROR(Node,
@@ -727,15 +728,6 @@ class TypeDecoder {
           flags.withConvention(FunctionMetadataConvention::CFunctionPointer);
       } else if (Node->getKind() == NodeKind::ThinFunctionType) {
         flags = flags.withConvention(FunctionMetadataConvention::Thin);
-      } else if (Node->getKind() == NodeKind::DifferentiableFunctionType ||
-               Node->getKind() ==
-                   NodeKind::EscapingDifferentiableFunctionType) {
-        flags = flags.withDifferentiabilityKind(
-            FunctionMetadataDifferentiabilityKind::Reverse);
-      } else if (Node->getKind() == NodeKind::LinearFunctionType ||
-                 Node->getKind() == NodeKind::EscapingLinearFunctionType) {
-        flags = flags.withDifferentiabilityKind(
-            FunctionMetadataDifferentiabilityKind::Linear);
       }
 
       unsigned firstChildIdx = 0;
@@ -766,8 +758,34 @@ class TypeDecoder {
         ++firstChildIdx;
       }
 
+      FunctionMetadataDifferentiabilityKind diffKind;
+      if (Node->getChild(firstChildIdx)->getKind() ==
+            NodeKind::DifferentiableFunctionType) {
+        auto mangledDiffKind = (MangledDifferentiabilityKind)
+            Node->getChild(firstChildIdx)->getIndex();
+        switch (mangledDiffKind) {
+        case MangledDifferentiabilityKind::NonDifferentiable:
+          assert(false && "Unexpected case NonDifferentiable");
+          break;
+        case MangledDifferentiabilityKind::Forward:
+          diffKind = FunctionMetadataDifferentiabilityKind::Forward;
+          break;
+        case MangledDifferentiabilityKind::Reverse:
+          diffKind = FunctionMetadataDifferentiabilityKind::Reverse;
+          break;
+        case MangledDifferentiabilityKind::Normal:
+          diffKind = FunctionMetadataDifferentiabilityKind::Normal;
+          break;
+        case MangledDifferentiabilityKind::Linear:
+          diffKind = FunctionMetadataDifferentiabilityKind::Linear;
+          break;
+        }
+        ++firstChildIdx;
+      }
+
       flags = flags.withConcurrent(isSendable)
-          .withAsync(isAsync).withThrows(isThrow);
+          .withAsync(isAsync).withThrows(isThrow)
+          .withDifferentiable(diffKind.isDifferentiable());
 
       if (Node->getNumChildren() < firstChildIdx + 2)
         return MAKE_NODE_TYPE_ERROR(Node,
@@ -786,16 +804,13 @@ class TypeDecoder {
               .withEscaping(
                           Node->getKind() == NodeKind::FunctionType ||
                           Node->getKind() == NodeKind::EscapingAutoClosureType ||
-                          Node->getKind() == NodeKind::EscapingObjCBlock ||
-                          Node->getKind() ==
-                              NodeKind::EscapingDifferentiableFunctionType ||
-                          Node->getKind() ==
-                              NodeKind::EscapingLinearFunctionType);
+                          Node->getKind() == NodeKind::EscapingObjCBlock);
 
       auto result = decodeMangledType(Node->getChild(firstChildIdx+1));
       if (result.isError())
         return result;
-      return Builder.createFunctionType(parameters, result.getType(), flags);
+      return Builder.createFunctionType(
+          parameters, result.getType(), flags, diffKind);
     }
     case NodeKind::ImplFunctionType: {
       auto calleeConvention = ImplParameterConvention::Direct_Unowned;

include/swift/Reflection/TypeRef.h

@@ -461,9 +461,11 @@ class FunctionTypeRef final : public TypeRef {
   std::vector<Param> Parameters;
   const TypeRef *Result;
   FunctionTypeFlags Flags;
+  FunctionMetadataDifferentiabilityKind DifferentiabilityKind;
 
   static TypeRefID Profile(const std::vector<Param> &Parameters,
-                           const TypeRef *Result, FunctionTypeFlags Flags) {
+                           const TypeRef *Result, FunctionTypeFlags Flags,
+                           FunctionMetadataDifferentiabilityKind DiffKind) {
     TypeRefID ID;
     for (const auto &Param : Parameters) {
       ID.addString(Param.getLabel().str());
@@ -472,20 +474,22 @@ class FunctionTypeRef final : public TypeRef {
     }
     ID.addPointer(Result);
     ID.addInteger(static_cast<uint64_t>(Flags.getIntValue()));
+    ID.addInteger(static_cast<uint64_t>(DiffKind.getIntValue()));
     return ID;
   }
 
 public:
   FunctionTypeRef(std::vector<Param> Params, const TypeRef *Result,
-                  FunctionTypeFlags Flags)
+                  FunctionTypeFlags Flags,
+                  FunctionMetadataDifferentiabilityKind DiffKind)
       : TypeRef(TypeRefKind::Function), Parameters(Params), Result(Result),
-        Flags(Flags) {}
+        Flags(Flags), DifferentiabilityKind(DiffKind) {}
 
   template <typename Allocator>
-  static const FunctionTypeRef *create(Allocator &A, std::vector<Param> Params,
-                                       const TypeRef *Result,
-                                       FunctionTypeFlags Flags) {
-    FIND_OR_CREATE_TYPEREF(A, FunctionTypeRef, Params, Result, Flags);
+  static const FunctionTypeRef *create(
+      Allocator &A, std::vector<Param> Params, const TypeRef *Result,
+      FunctionTypeFlags Flags, FunctionMetadataDifferentiabilityKind DiffKind) {
+    FIND_OR_CREATE_TYPEREF(A, FunctionTypeRef, Params, Result, Flags, DiffKind);
   }
 
   const std::vector<Param> &getParameters() const { return Parameters; };
@@ -498,6 +502,10 @@ class FunctionTypeRef final : public TypeRef {
     return Flags;
   }
 
+  FunctionMetadataDifferentiabilityKind getDifferentiabilityKind() const {
+    return DifferentiabilityKind;
+  }
+
   static bool classof(const TypeRef *TR) {
     return TR->getKind() == TypeRefKind::Function;
   }

include/swift/Reflection/TypeRefBuilder.h

@@ -414,8 +414,9 @@ class TypeRefBuilder {
 
   const FunctionTypeRef *createFunctionType(
       llvm::ArrayRef<remote::FunctionParam<const TypeRef *>> params,
-      const TypeRef *result, FunctionTypeFlags flags) {
-    return FunctionTypeRef::create(*this, params, result, flags);
+      const TypeRef *result, FunctionTypeFlags flags,
+      FunctionMetadataDifferentiabilityKind diffKind) {
+    return FunctionTypeRef::create(*this, params, result, flags, diffKind);
   }
   using BuiltSubstitution = std::pair<const TypeRef *, const TypeRef *>;
   using BuiltRequirement = TypeRefRequirement;
@@ -454,9 +455,29 @@ class TypeRefBuilder {
 
     funcFlags = funcFlags.withConcurrent(flags.isSendable());
     funcFlags = funcFlags.withAsync(flags.isAsync());
+    funcFlags = funcFlags.withDifferentiable(flags.isDifferentiable());
+
+    FunctionMetadataDifferentiabilityKind diffKind;
+    switch (flags.getDifferentiabilityKind()) {
+    case ImplFunctionDifferentiabilityKind::NonDifferentiable:
+      diffKind = FunctionMetadataDifferentiabilityKind::NonDifferentiable;
+      break;
+    case ImplFunctionDifferentiabilityKind::Forward:
+      diffKind = FunctionMetadataDifferentiabilityKind::Forward;
+      break;
+    case ImplFunctionDifferentiabilityKind::Reverse:
+      diffKind = FunctionMetadataDifferentiabilityKind::Reverse;
+      break;
+    case ImplFunctionDifferentiabilityKind::Normal:
+      diffKind = FunctionMetadataDifferentiabilityKind::Normal;
+      break;
+    case ImplFunctionDifferentiabilityKind::Linear:
+      diffKind = FunctionMetadataDifferentiabilityKind::Linear;
+      break;
+    }
 
     auto result = createTupleType({}, "");
-    return FunctionTypeRef::create(*this, {}, result, funcFlags);
+    return FunctionTypeRef::create(*this, {}, result, funcFlags, diffKind);
   }
 
   const ProtocolCompositionTypeRef *