Merge branch 'tensorflow' of github.com:apple/swift into tensorflow-merge

```
Failing Tests (3):
    Swift(macosx-x86_64) :: AutoDiff/class_method_thunk/main.swift
    Swift(macosx-x86_64) :: AutoDiff/differentiable_attr_cross_module/main.swift
    Swift(macosx-x86_64) :: AutoDiff/silgen_thunking/main.swift
```

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1541,6 +1541,8 @@ class DifferentiableAttr final
                                     ParsedAutoDiffParameter> {
   friend TrailingObjects;
 
+  /// The declaration on which the `@differentiable` attribute is declared.
+  Decl *OriginalDeclaration = nullptr;
   /// Whether this function is linear.
   bool Linear;
   /// The number of parsed parameters specified in 'wrt:'.
@@ -1573,7 +1575,7 @@ class DifferentiableAttr final
                               Optional<DeclNameWithLoc> vjp,
                               TrailingWhereClause *clause);
 
-  explicit DifferentiableAttr(ASTContext &context, bool implicit,
+  explicit DifferentiableAttr(Decl *original, bool implicit,
                               SourceLoc atLoc, SourceRange baseRange,
                               bool linear, IndexSubset *indices,
                               Optional<DeclNameWithLoc> jvp,
@@ -1589,13 +1591,16 @@ class DifferentiableAttr final
                                     Optional<DeclNameWithLoc> vjp,
                                     TrailingWhereClause *clause);
 
-  static DifferentiableAttr *create(ASTContext &context, bool implicit,
+  static DifferentiableAttr *create(Decl *original, bool implicit,
                                     SourceLoc atLoc, SourceRange baseRange,
                                     bool linear, IndexSubset *indices,
                                     Optional<DeclNameWithLoc> jvp,
                                     Optional<DeclNameWithLoc> vjp,
                                     GenericSignature derivativeGenSig);
 
+  Decl *getOriginalDeclaration() const { return OriginalDeclaration; }
+  void setOriginalDeclaration(Decl *decl);
+
   /// Get the optional 'jvp:' function name and location.
   /// Use this instead of `getJVPFunction` to check whether the attribute has a
   /// registered JVP.

include/swift/AST/DiagnosticsSema.def

@@ -2974,8 +2974,8 @@ ERROR(transpose_params_clause_param_not_differentiable,none,
       "'Differentiable' and where '%0 == %0.TangentVector'", (StringRef))
 ERROR(transposing_attr_overload_not_found,none,
       "could not find function %0 with expected type %1", (DeclName, Type))
-ERROR(transposing_attr_cant_use_named_wrt_params,none,
-      "cannot use named wrt parameters in '@transposing' attribute, found %0",
+ERROR(transposing_attr_cannot_use_named_wrt_params,none,
+      "cannot use named 'wrt' parameters in '@transposing' attribute, found %0",
       (Identifier))
 ERROR(transposing_attr_result_value_not_differentiable,none,
       "'@transposing' attribute requires original function result to "

include/swift/AST/Types.h

@@ -3139,12 +3139,11 @@ class AnyFunctionType : public TypeBase {
   /// Given the type of an autodiff derivative function, returns the
   /// corresponding original function type.
   AnyFunctionType *getAutoDiffOriginalFunctionType();
-  
+
   /// Given the type of a transposing derivative function, returns the
   /// corresponding original function type.
   AnyFunctionType *
-  getTransposeOriginalFunctionType(TransposingAttr *attr,
-                                   IndexSubset *wrtParamIndices,
+  getTransposeOriginalFunctionType(IndexSubset *wrtParamIndices,
                                    bool wrtSelf);
 
   AnyFunctionType *getWithoutDifferentiability() const;

lib/AST/Attr.cpp

@@ -1454,16 +1454,16 @@ DifferentiableAttr::DifferentiableAttr(ASTContext &context, bool implicit,
             getTrailingObjects<ParsedAutoDiffParameter>());
 }
 
-DifferentiableAttr::DifferentiableAttr(ASTContext &context, bool implicit,
+DifferentiableAttr::DifferentiableAttr(Decl *original, bool implicit,
                                        SourceLoc atLoc, SourceRange baseRange,
-                                       bool linear,
-                                       IndexSubset *indices,
+                                       bool linear, IndexSubset *indices,
                                        Optional<DeclNameWithLoc> jvp,
                                        Optional<DeclNameWithLoc> vjp,
                                        GenericSignature derivativeGenSig)
     : DeclAttribute(DAK_Differentiable, atLoc, baseRange, implicit),
       Linear(linear), JVP(std::move(jvp)), VJP(std::move(vjp)),
       ParameterIndices(indices) {
+  setOriginalDeclaration(original);
   setDerivativeGenericSignature(derivativeGenSig);
 }
 
@@ -1483,19 +1483,26 @@ DifferentiableAttr::create(ASTContext &context, bool implicit,
 }
 
 DifferentiableAttr *
-DifferentiableAttr::create(ASTContext &context, bool implicit,
-                           SourceLoc atLoc, SourceRange baseRange,
-                           bool linear, IndexSubset *indices,
-                           Optional<DeclNameWithLoc> jvp,
+DifferentiableAttr::create(Decl *original, bool implicit, SourceLoc atLoc,
+                           SourceRange baseRange, bool linear,
+                           IndexSubset *indices, Optional<DeclNameWithLoc> jvp,
                            Optional<DeclNameWithLoc> vjp,
                            GenericSignature derivativeGenSig) {
-  void *mem = context.Allocate(sizeof(DifferentiableAttr),
-                               alignof(DifferentiableAttr));
-  return new (mem) DifferentiableAttr(context, implicit, atLoc, baseRange,
+  auto &ctx = original->getASTContext();
+  void *mem = ctx.Allocate(sizeof(DifferentiableAttr),
+                           alignof(DifferentiableAttr));
+  return new (mem) DifferentiableAttr(original, implicit, atLoc, baseRange,
                                       linear, indices, std::move(jvp),
                                       std::move(vjp), derivativeGenSig);
 }
 
+void DifferentiableAttr::setOriginalDeclaration(Decl *decl) {
+  assert(decl && "Original declaration must be non-null");
+  assert(!OriginalDeclaration &&
+         "Original declaration cannot have already been set");
+  OriginalDeclaration = decl;
+}
+
 void DifferentiableAttr::setJVPFunction(FuncDecl *decl) {
   JVPFunction = decl;
   if (decl && !JVP)

lib/AST/Type.cpp

@@ -4867,25 +4867,24 @@ makeFunctionType(ArrayRef<AnyFunctionType::Param> params, Type retTy,
 // Compute the original function type corresponding to the given transpose
 // function type.
 AnyFunctionType *AnyFunctionType::getTransposeOriginalFunctionType(
-    TransposingAttr *attr, IndexSubset *wrtParamIndices, bool wrtSelf) {
+    IndexSubset *wrtParamIndices, bool wrtSelf) {
   unsigned transposeParamsIndex = 0;
   bool isCurried = getResult()->is<AnyFunctionType>();
-  
+
   // Get the original function's result.
   auto transposeParams = getParams();
   auto transposeResult = getResult();
   if (isCurried) {
-    auto method =
-        getAs<AnyFunctionType>()->getResult()->getAs<AnyFunctionType>();
-    transposeParams = method->getParams();
-    transposeResult = method->getResult();
+    auto methodType = getResult()->castTo<AnyFunctionType>();
+    transposeParams = methodType->getParams();
+    transposeResult = methodType->getResult();
   }
-      
+
   Type originalResult;
   if (isCurried) {
     // If it's curried, then the first parameter in the curried type, which is
     // the 'Self' type, is the original result (no matter if we are
-    // differentiating WRT self or aren't).
+    // transposing wrt self or not).
     originalResult = getParams().front().getPlainType();
   } else {
     // If it's not curried, the last parameter, the tangent, is always the
@@ -4895,22 +4894,21 @@ AnyFunctionType *AnyFunctionType::getTransposeOriginalFunctionType(
   }
   assert(originalResult);
 
-  auto wrtParams = attr->getParsedParameters();
   SmallVector<TupleTypeElt, 4> transposeResultTypes;
   // Return type of '@transposing' function can have single type or tuples
   // of types.
-  if (auto t = transposeResult->getAs<TupleType>()) {
-    transposeResultTypes.append(t->getElements().begin(),
-                                t->getElements().end());
+  if (auto transposeResultTupleType = transposeResult->getAs<TupleType>()) {
+    transposeResultTypes.append(transposeResultTupleType->getElements().begin(),
+                                transposeResultTupleType->getElements().end());
   } else {
     transposeResultTypes.push_back(transposeResult);
   }
   assert(!transposeResultTypes.empty());
 
-  // If the function is curried and is transposing WRT 'self', then grab
+  // If the function is curried and is transposing wrt 'self', then grab
   // the type from the result list (guaranteed to be the first since 'self'
-  // is first in WRT list) and remove it. If it's still curried but not
-  // transposing WRT 'self', then the 'Self' type is the first parameter
+  // is first in wrt list) and remove it. If it is still curried but not
+  // transposing wrt 'self', then the 'Self' type is the first parameter
   // in the method.
   unsigned transposeResultTypesIndex = 0;
   Type selfType;
@@ -4923,21 +4921,21 @@ AnyFunctionType *AnyFunctionType::getTransposeOriginalFunctionType(
   }
 
   SmallVector<AnyFunctionType::Param, 8> originalParams;
-  unsigned numberOriginalParameters =
-      transposeParams.size() + wrtParams.size() - 1;
-  for (auto i : range(numberOriginalParameters)) {
+  unsigned originalParameterCount =
+      transposeParams.size() + wrtParamIndices->getNumIndices() - 1;
+  for (auto i : range(originalParameterCount)) {
     // Need to check if it is the 'self' param since we handle it differently
     // above.
-    bool lookingAtSelf = (i == (wrtParamIndices->getCapacity() - 1)) && wrtSelf;
-    bool isWrt = wrtParamIndices->contains(i);
-    if (isWrt && !lookingAtSelf) {
-      // If in WRT list, the item in the result tuple must be a parameter in the
+    bool lookingAtSelf = (i == wrtParamIndices->getCapacity() - 1) && wrtSelf;
+    if (wrtParamIndices->contains(i) && !lookingAtSelf) {
+      // If in wrt list, the item in the result tuple must be a parameter in the
       // original function.
-      auto resultType = transposeResultTypes[transposeResultTypesIndex].getType();
+      auto resultType =
+          transposeResultTypes[transposeResultTypesIndex].getType();
       originalParams.push_back(AnyFunctionType::Param(resultType));
       transposeResultTypesIndex++;
     } else {
-      // Else if not in the WRT list, the parameter in the transposing function
+      // Else if not in the wrt list, the parameter in the transposing function
       // is a parameter in the original function.
       originalParams.push_back(transposeParams[transposeParamsIndex]);
       transposeParamsIndex++;

lib/Parse/ParseDecl.cpp

@@ -3384,6 +3384,13 @@ void Parser::delayParseFromBeginningToHere(ParserPosition BeginParserPosition,
     consumeToken();
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+static void setOriginalFunctionInDifferentiableAttributes(
+    DeclAttributes Attributes, Decl *D) {
+  for (auto *attr : Attributes.getAttributes<DifferentiableAttr>())
+    const_cast<DifferentiableAttr *>(attr)->setOriginalDeclaration(D);
+}
+
 /// Parse a single syntactic declaration and return a list of decl
 /// ASTs.  This can return multiple results for var decls that bind to multiple
 /// values, structs that define a struct decl and a constructor, etc.
@@ -3788,6 +3795,7 @@ Parser::parseDecl(ParseDeclOptions Flags,
     Decl *D = DeclResult.get();
     if (!declWasHandledAlready(D)) {
       Handler(D);
+      // SWIFT_ENABLE_TENSORFLOW
       if (auto FD = dyn_cast<FuncDecl>(D)) {
         if (auto attr = D->getAttrs().getAttribute<QuotedAttr>()) {
           // TODO(TF-718): Properly mangle names for quote decls.
@@ -3825,7 +3833,11 @@ Parser::parseDecl(ParseDeclOptions Flags,
           Handler(quoteDecl);
         }
       }
+      // SWIFT_ENABLE_TENSORFLOW END
     }
+    // SWIFT_ENABLE_TENSORFLOW
+    setOriginalFunctionInDifferentiableAttributes(D->getAttrs(), D);
+    // SWIFT_ENABLE_TENSORFLOW END
   }
 
   if (!DeclResult.isParseError()) {
@@ -5579,6 +5591,12 @@ Parser::parseDeclVarGetSet(Pattern *pattern, ParseDeclOptions Flags,
 
   accessors.record(*this, PrimaryVar, Invalid);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  for (auto *accessor : accessors.Accessors)
+    setOriginalFunctionInDifferentiableAttributes(accessor->getAttrs(),
+                                                  accessor);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   return makeParserResult(PrimaryVar);
 }
 
@@ -5833,6 +5851,9 @@ Parser::parseDeclVar(ParseDeclOptions Flags,
     pattern->forEachVariable([&](VarDecl *VD) {
       VD->setStatic(StaticLoc.isValid());
       VD->getAttrs() = Attributes;
+      // SWIFT_ENABLE_TENSORFLOW
+      setOriginalFunctionInDifferentiableAttributes(Attributes, VD);
+      // SWIFT_ENABLE_TENSORFLOW END
       setLocalDiscriminator(VD);
       Decls.push_back(VD);
       if (hasOpaqueReturnTy && sf && !InInactiveClauseEnvironment) {
@@ -7087,6 +7108,12 @@ Parser::parseDeclSubscript(SourceLoc StaticLoc,
 
   accessors.record(*this, Subscript, (Invalid || !Status.isSuccess()));
 
+  // SWIFT_ENABLE_TENSORFLOW
+  for (auto *accessor : accessors.Accessors)
+    setOriginalFunctionInDifferentiableAttributes(accessor->getAttrs(),
+                                                  accessor);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   // No need to setLocalDiscriminator because subscripts cannot
   // validly appear outside of type decls.
   return makeParserResult(Status, Subscript);

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -628,8 +628,9 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
       if (auto *extDecl = dyn_cast<ExtensionDecl>(parentDC->getAsDecl()))
         derivativeGenSig = extDecl->getGenericSignature();
       auto *diffableAttr = DifferentiableAttr::create(
-          C, /*implicit*/ true, SourceLoc(), SourceLoc(),
-          /*linear*/ false, {}, None, None, derivativeGenSig);
+          member->getAccessor(AccessorKind::Get), /*implicit*/ true,
+          SourceLoc(), SourceLoc(), /*linear*/ false, {}, None, None,
+          derivativeGenSig);
       member->getAttrs().add(diffableAttr);
       // Set getter `@differentiable` attribute parameter indices.
       diffableAttr->setParameterIndices(IndexSubset::get(C, 1, {0}));