Merge pull request swiftlang#35847 from rxwei/72666310-mangle-ad-thunks

Author: swift-ci

docs/ABI/Mangling.rst

@@ -220,7 +220,6 @@ types where the metadata itself has unknown layout.)
   global ::= global 'Tm'                 // merged function
   global ::= entity                      // some identifiable thing
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
-  global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
   global ::= impl-function-type type 'Tz'     // objc-to-swift-async completion handler block implementation
   global ::= impl-function-type type 'TZ'     // objc-to-swift-async completion handler block implementation (predefined by runtime)
   global ::= from-type to-type self-type generic-signature? 'Ty'  // reabstraction thunk with dynamic 'Self' capture
@@ -230,6 +229,9 @@ types where the metadata itself has unknown layout.)
   global ::= type generic-signature 'TH' // key path equality
   global ::= type generic-signature 'Th' // key path hasher
   global ::= global generic-signature? 'TJ' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' // autodiff function
+  global ::= from-type to-type 'TJO' AUTODIFF-FUNCTION-KIND // autodiff self-reordering reabstraction thunk
+  global ::= from-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff linear map subset parameters thunk
+  global ::= global to-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff derivative function subset parameters thunk
 
   global ::= protocol 'TL'               // protocol requirements base descriptor
   global ::= assoc-type-name 'Tl'        // associated type descriptor

include/swift/AST/ASTMangler.h

@@ -192,6 +192,15 @@ class ASTMangler : public Mangler {
                                       AutoDiffLinearMapKind kind,
                                       AutoDiffConfig config);
 
+  /// Mangle the linear map self parameter reordering thunk the given:
+  /// - Mangled original function declaration.
+  /// - Linear map kind.
+  /// - Derivative function configuration: parameter/result indices and
+  ///   derivative generic signature.
+  std::string mangleAutoDiffSelfReorderingReabstractionThunk(
+      CanType fromType, CanType toType, GenericSignature signature,
+      AutoDiffLinearMapKind linearMapKind);
+
   /// Mangle the AutoDiff generated declaration for the given:
   /// - Generated declaration kind: linear map struct or branching trace enum.
   /// - Mangled original function name.

include/swift/Demangling/DemangleNodes.def

@@ -310,6 +310,8 @@ NODE(CanonicalPrespecializedGenericTypeCachingOnceToken)
 NODE(AsyncFunctionPointer)
 CONTEXT_NODE(AutoDiffFunction)
 NODE(AutoDiffFunctionKind)
+NODE(AutoDiffSelfReorderingReabstractionThunk)
+NODE(AutoDiffSubsetParametersThunk)
 NODE(IndexSubset)
 
 #undef CONTEXT_NODE

include/swift/Demangling/Demangler.h

@@ -569,7 +569,10 @@ class Demangler : public NodeFactory {
 
   NodePointer demangleTypeMangling();
   NodePointer demangleSymbolicReference(unsigned char rawKind);
+  NodePointer demangleAutoDiffFunction();
   NodePointer demangleAutoDiffFunctionKind();
+  NodePointer demangleAutoDiffSubsetParametersThunk();
+  NodePointer demangleAutoDiffSelfReorderingReabstractionThunk();
   NodePointer demangleIndexSubset();
 
   bool demangleBoundGenerics(Vector<NodePointer> &TypeListList,

include/swift/SILOptimizer/Utils/DifferentiationMangler.h

@@ -27,17 +27,29 @@ class DifferentiationMangler : public ASTMangler {
 public:
   DifferentiationMangler() {}
   /// Returns the mangled name for a differentiation function of the given kind.
-  std::string mangle(SILFunction *originalFunction,
-                     Demangle::AutoDiffFunctionKind kind,
-                     AutoDiffConfig config);
+  std::string mangleAutoDiffFunction(StringRef originalName,
+                                     Demangle::AutoDiffFunctionKind kind,
+                                     AutoDiffConfig config);
   /// Returns the mangled name for a derivative function of the given kind.
-  std::string mangleDerivativeFunction(SILFunction *originalFunction,
+  std::string mangleDerivativeFunction(StringRef originalName,
                                        AutoDiffDerivativeFunctionKind kind,
                                        AutoDiffConfig config);
   /// Returns the mangled name for a linear map of the given kind.
-  std::string mangleLinearMap(SILFunction *originalFunction,
+  std::string mangleLinearMap(StringRef originalName,
                               AutoDiffLinearMapKind kind,
                               AutoDiffConfig config);
+  /// Returns the mangled name for a derivative function subset parameters
+  /// thunk.
+  std::string mangleDerivativeFunctionSubsetParametersThunk(
+      StringRef originalName, CanType toType,
+      AutoDiffDerivativeFunctionKind linearMapKind,
+      IndexSubset *fromParamIndices, IndexSubset *fromResultIndices,
+      IndexSubset *toParamIndices);
+  /// Returns the mangled name for a linear map subset parameters thunk.
+  std::string mangleLinearMapSubsetParametersThunk(
+      CanType fromType, AutoDiffLinearMapKind linearMapKind,
+      IndexSubset *fromParamIndices, IndexSubset *fromResultIndices,
+      IndexSubset *toParamIndices);
 };
 
 } // end namespace Mangle

lib/AST/ASTMangler.cpp

@@ -448,6 +448,19 @@ void ASTMangler::appendAutoDiffFunctionParts(char functionKindCode,
   appendOperator("r");
 }
 
+std::string ASTMangler::mangleAutoDiffSelfReorderingReabstractionThunk(
+    CanType fromType, CanType toType, GenericSignature signature,
+    AutoDiffLinearMapKind linearMapKind) {
+  beginMangling();
+  appendType(fromType);
+  appendType(toType);
+  if (signature)
+    appendGenericSignature(signature);
+  auto kindCode = (char)getAutoDiffFunctionKind(linearMapKind);
+  appendOperator("TJO", StringRef(&kindCode, 1));
+  return finalize();
+}
+
 /// Mangle the index subset.
 void ASTMangler::appendIndexSubset(IndexSubset *indices) {
   Buffer << indices->getString();

lib/Demangling/Demangler.cpp

@@ -2498,34 +2498,74 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
       return createNode(Node::Kind::OutlinedBridgedMethod, Params);
     }
     case 'u': return createNode(Node::Kind::AsyncFunctionPointer);
-    case 'J': {
-      auto result = createNode(Node::Kind::AutoDiffFunction);
-      auto optionalGenSig = popNode(Node::Kind::DependentGenericSignature);
-      auto original = popNode();
-      result = addChild(result, original);
-      addChild(result, optionalGenSig);
-      auto kind = demangleAutoDiffFunctionKind();
-      if (!kind)
-        return nullptr;
-      result = addChild(result, kind);
-      result = addChild(result, demangleIndexSubset());
-      if (!nextIf('p')) return nullptr;
-      result = addChild(result, demangleIndexSubset());
-      if (!nextIf('r')) return nullptr;
-      return result;
-    }
+    case 'J':
+      switch (peekChar()) {
+      case 'S':
+        return demangleAutoDiffSubsetParametersThunk();
+      case 'O':
+        return demangleAutoDiffSelfReorderingReabstractionThunk();
+      }
+      return demangleAutoDiffFunction();
     default:
       return nullptr;
   }
 }
 
+NodePointer Demangler::demangleAutoDiffFunction() {
+  auto result = createNode(Node::Kind::AutoDiffFunction);
+  while (auto *originalNode = popNode())
+    result = addChild(result, originalNode);
+  result->reverseChildren();
+  auto kind = demangleAutoDiffFunctionKind();
+  result = addChild(result, kind);
+  result = addChild(result, demangleIndexSubset());
+  if (!nextIf('p'))
+    return nullptr;
+  result = addChild(result, demangleIndexSubset());
+  if (!nextIf('r'))
+    return nullptr;
+  return result;
+}
+
 NodePointer Demangler::demangleAutoDiffFunctionKind() {
   char kind = nextChar();
   if (kind != 'f' && kind != 'r' && kind != 'd' && kind != 'p')
     return nullptr;
   return createNode(Node::Kind::AutoDiffFunctionKind, kind);
 }
 
+NodePointer Demangler::demangleAutoDiffSubsetParametersThunk() {
+  nextChar();
+  auto result = createNode(Node::Kind::AutoDiffSubsetParametersThunk);
+  while (auto *node = popNode())
+    result = addChild(result, node);
+  result->reverseChildren();
+  auto kind = demangleAutoDiffFunctionKind();
+  result = addChild(result, kind);
+  result = addChild(result, demangleIndexSubset());
+  if (!nextIf('p'))
+    return nullptr;
+  result = addChild(result, demangleIndexSubset());
+  if (!nextIf('r'))
+    return nullptr;
+  result = addChild(result, demangleIndexSubset());
+  if (!nextIf('P'))
+    return nullptr;
+  return result;
+}
+
+NodePointer Demangler::demangleAutoDiffSelfReorderingReabstractionThunk() {
+  nextChar();
+  auto result = createNode(
+      Node::Kind::AutoDiffSelfReorderingReabstractionThunk);
+  addChild(result, popNode(Node::Kind::DependentGenericSignature));
+  result = addChild(result, popNode(Node::Kind::Type));
+  result = addChild(result, popNode(Node::Kind::Type));
+  result->reverseChildren();
+  result = addChild(result, demangleAutoDiffFunctionKind());
+  return result;
+}
+
 NodePointer Demangler::demangleIndexSubset() {
   std::string str;
   for (auto c = peekChar(); c == 'S' || c == 'U'; c = peekChar()) {

lib/Demangling/NodePrinter.cpp

@@ -566,6 +566,8 @@ class NodePrinter {
     case Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken:
     case Node::Kind::AsyncFunctionPointer:
     case Node::Kind::AutoDiffFunction:
+    case Node::Kind::AutoDiffSelfReorderingReabstractionThunk:
+    case Node::Kind::AutoDiffSubsetParametersThunk:
     case Node::Kind::AutoDiffFunctionKind:
     case Node::Kind::IndexSubset:
       return false;
@@ -1738,17 +1740,29 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     return nullptr;
   }
   case Node::Kind::AutoDiffFunction: {
-    auto childIt = Node->begin();
-    auto original = *childIt++;
-    NodePointer optionalGenSig =
-        (*childIt)->getKind() == Node::Kind::DependentGenericSignature
-            ? *childIt++ : nullptr;
-    auto kind = *childIt++;
-    auto paramIndices = *childIt++;
-    auto resultIndices = *childIt++;
+    unsigned prefixEndIndex = 0;
+    while (prefixEndIndex != Node->getNumChildren() &&
+           Node->getChild(prefixEndIndex)->getKind()
+              != Node::Kind::AutoDiffFunctionKind)
+      ++prefixEndIndex;
+    auto kind = Node->getChild(prefixEndIndex);
+    auto paramIndices = Node->getChild(prefixEndIndex + 1);
+    auto resultIndices = Node->getChild(prefixEndIndex + 2);
     print(kind);
     Printer << " of ";
-    print(original);
+    NodePointer optionalGenSig = nullptr;
+    for (unsigned i = 0; i < prefixEndIndex; ++i) {
+      // The last node may be a generic signature. If so, print it later.
+      if (i == prefixEndIndex - 1 &&
+          Node->getChild(i)->getKind()
+              == Node::Kind::DependentGenericSignature) {
+        optionalGenSig = Node->getChild(i);
+        break;
+      }
+      print(Node->getChild(i));
+    }
+    if (Options.ShortenThunk)
+      return nullptr;
     Printer << " with respect to parameters ";
     print(paramIndices);
     Printer << " and results ";
@@ -1759,6 +1773,61 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     }
     return nullptr;
   }
+  case Node::Kind::AutoDiffSelfReorderingReabstractionThunk: {
+    Printer << "autodiff self-reordering reabstraction thunk ";
+    auto childIt = Node->begin();
+    auto fromType = *childIt++;
+    auto toType = *childIt++;
+    if (Options.ShortenThunk) {
+      Printer << "for ";
+      print(fromType);
+      return nullptr;
+    }
+    NodePointer optionalGenSig =
+        (*childIt)->getKind() == Node::Kind::DependentGenericSignature
+            ? *childIt++ : nullptr;
+    Printer << "for ";
+    print(*childIt++); // kind
+    if (optionalGenSig) {
+      print(optionalGenSig);
+      Printer << ' ';
+    }
+    Printer << " from ";
+    print(fromType);
+    Printer << " to ";
+    print(toType);
+    return nullptr;
+  }
+  case Node::Kind::AutoDiffSubsetParametersThunk: {
+    Printer << "autodiff subset parameters thunk for ";
+    auto currentIndex = Node->getNumChildren() - 1;
+    auto toParamIndices = Node->getChild(currentIndex--);
+    auto resultIndices = Node->getChild(currentIndex--);
+    auto paramIndices = Node->getChild(currentIndex--);
+    auto kind = Node->getChild(currentIndex--);
+    print(kind);
+    Printer << " from ";
+    // Print the "from" thing.
+    if (currentIndex == 0) {
+      print(Node->getFirstChild()); // the "from" type
+    } else {
+      for (unsigned i = 0; i < currentIndex; ++i) // the "from" global
+        print(Node->getChild(i));
+    }
+    if (Options.ShortenThunk)
+      return nullptr;
+    Printer << " with respect to parameters ";
+    print(paramIndices);
+    Printer << " and results ";
+    print(resultIndices);
+    Printer << " to parameters ";
+    print(toParamIndices);
+    if (currentIndex > 0) {
+      Printer << " of type ";
+      print(Node->getChild(currentIndex)); // "to" type
+    }
+    return nullptr;
+  }
   case Node::Kind::AutoDiffFunctionKind: {
     auto kind = (AutoDiffFunctionKind)Node->getIndex();
     switch (kind) {

lib/Demangling/OldRemangler.cpp

@@ -748,6 +748,14 @@ void Remangler::mangleAutoDiffFunction(Node *node, EntityContext &ctx) {
   Buffer << "<autodiff-function>";
 }
 
+void Remangler::mangleAutoDiffSelfReorderingReabstractionThunk(Node *node) {
+  Buffer << "<autodiff-self-reordering-reabstraction-thunk>";
+}
+
+void Remangler::mangleAutoDiffSubsetParametersThunk(Node *node) {
+  Buffer << "<autodiff-subset-parameters-thunk>";
+}
+
 void Remangler::mangleAutoDiffFunctionKind(Node *node) {
   Buffer << "<autodiff-function-kind>";
 }

lib/Demangling/Remangler.cpp

@@ -2116,15 +2116,40 @@ void Remangler::mangleReabstractionThunkHelperWithSelf(Node *node) {
 
 void Remangler::mangleAutoDiffFunction(Node *node) {
   auto childIt = node->begin();
-  mangle(*childIt++); // original
+  while (childIt != node->end() &&
+         (*childIt)->getKind() != Node::Kind::AutoDiffFunctionKind)
+    mangle(*childIt++);
+  Buffer << "TJ";
+  mangle(*childIt++); // kind
+  mangle(*childIt++); // parameter indices
+  Buffer << 'p';
+  mangle(*childIt++); // result indices
+  Buffer << 'r';
+}
+
+void Remangler::mangleAutoDiffSelfReorderingReabstractionThunk(Node *node) {
+  auto childIt = node->begin();
+  mangle(*childIt++); // from type
+  mangle(*childIt++); // to type
   if ((*childIt)->getKind() == Node::Kind::DependentGenericSignature)
     mangleDependentGenericSignature(*childIt++);
-  Buffer << "TJ";
+  Buffer << "TJO";
+  mangle(*childIt++); // kind
+}
+
+void Remangler::mangleAutoDiffSubsetParametersThunk(Node *node) {
+  auto childIt = node->begin();
+  while (childIt != node->end() &&
+         (*childIt)->getKind() != Node::Kind::AutoDiffFunctionKind)
+    mangle(*childIt++);
+  Buffer << "TJS";
   mangle(*childIt++); // kind
   mangle(*childIt++); // parameter indices
   Buffer << 'p';
   mangle(*childIt++); // result indices
   Buffer << 'r';
+  mangle(*childIt++); // to parameter indices
+  Buffer << 'P';
 }
 
 void Remangler::mangleAutoDiffFunctionKind(Node *node) {