Merge pull request #1442 from swiftwasm/master

[pull] swiftwasm from master

Author: pull[bot]

include/swift/AST/Attr.h

@@ -1715,13 +1715,6 @@ class OriginallyDefinedInAttr: public DeclAttribute {
   }
 };
 
-/// A declaration name with location.
-struct DeclNameRefWithLoc {
-  DeclNameRef Name;
-  DeclNameLoc Loc;
-  Optional<AccessorKind> AccessorKind;
-};
-
 /// Attribute that marks a function as differentiable.
 ///
 /// Examples:
@@ -1847,6 +1840,18 @@ class DifferentiableAttr final
   }
 };
 
+/// A declaration name with location.
+struct DeclNameRefWithLoc {
+  /// The declaration name.
+  DeclNameRef Name;
+  /// The declaration name location.
+  DeclNameLoc Loc;
+  /// An optional accessor kind.
+  Optional<AccessorKind> AccessorKind;
+
+  void print(ASTPrinter &Printer) const;
+};
+
 /// The `@derivative(of:)` attribute registers a function as a derivative of
 /// another function-like declaration: a 'func', 'init', 'subscript', or 'var'
 /// computed property declaration.

lib/AST/Attr.cpp

@@ -1052,7 +1052,9 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@derivative");
     Printer << "(of: ";
     auto *attr = cast<DerivativeAttr>(this);
-    Printer << attr->getOriginalFunctionName().Name;
+    if (auto *baseType = attr->getBaseTypeRepr())
+      baseType->print(Printer, Options);
+    attr->getOriginalFunctionName().print(Printer);
     auto *derivative = cast<AbstractFunctionDecl>(D);
     auto diffParamsString = getDifferentiationParametersClauseString(
         derivative, attr->getParameterIndices(), attr->getParsedParameters(),
@@ -1067,7 +1069,9 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@transpose");
     Printer << "(of: ";
     auto *attr = cast<TransposeAttr>(this);
-    Printer << attr->getOriginalFunctionName().Name;
+    if (auto *baseType = attr->getBaseTypeRepr())
+      baseType->print(Printer, Options);
+    attr->getOriginalFunctionName().print(Printer);
     auto *transpose = cast<AbstractFunctionDecl>(D);
     auto transParamsString = getDifferentiationParametersClauseString(
         transpose, attr->getParameterIndices(), attr->getParsedParameters(),
@@ -1719,6 +1723,12 @@ GenericEnvironment *DifferentiableAttr::getDerivativeGenericEnvironment(
   return original->getGenericEnvironment();
 }
 
+void DeclNameRefWithLoc::print(ASTPrinter &Printer) const {
+  Printer << Name;
+  if (AccessorKind)
+    Printer << '.' << getAccessorLabel(*AccessorKind);
+}
+
 void DifferentiableAttr::print(llvm::raw_ostream &OS, const Decl *D,
                                bool omitWrtClause) const {
   StreamPrinter P(OS);

lib/AST/AutoDiff.cpp

@@ -422,6 +422,14 @@ void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
   }
 }
 
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const DeclNameRefWithLoc &name) {
+  os << name.Name;
+  if (auto accessorKind = name.AccessorKind)
+    os << '.' << getAccessorLabel(*accessorKind);
+  return os;
+}
+
 bool swift::operator==(const TangentPropertyInfo::Error &lhs,
                        const TangentPropertyInfo::Error &rhs) {
   if (lhs.kind != rhs.kind)

lib/SILOptimizer/Analysis/EscapeAnalysis.cpp

@@ -2079,6 +2079,10 @@ void EscapeAnalysis::analyzeInstruction(SILInstruction *I,
       return;
   }
 
+  // Incidental uses produce no values and have no effect on their operands.
+  if (isIncidentalUse(I))
+    return;
+
   // Instructions which return the address of non-writable memory cannot have
   // an effect on escaping.
   if (isNonWritableMemoryAddress(I))

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -457,10 +457,6 @@ static void addPerfEarlyModulePassPipeline(SILPassPipelinePlan &P) {
   // optimization.
   P.addGlobalOpt();
 
-  // We earlier eliminated ownership if we are not compiling the stdlib. Now
-  // handle the stdlib functions.
-  P.addNonTransparentFunctionOwnershipModelEliminator();
-
   // Add the outliner pass (Osize).
   P.addOutliner();
 
@@ -485,6 +481,11 @@ static void addHighLevelFunctionPipeline(SILPassPipelinePlan &P) {
   P.startPipeline("HighLevel,Function+EarlyLoopOpt");
   // FIXME: update EagerSpecializer to be a function pass!
   P.addEagerSpecializer();
+
+  // We earlier eliminated ownership if we are not compiling the stdlib. Now
+  // handle the stdlib functions.
+  P.addNonTransparentFunctionOwnershipModelEliminator();
+
   addFunctionPasses(P, OptimizationLevelKind::HighLevel);
 
   addHighLevelLoopOptPasses(P);
@@ -714,6 +715,8 @@ SILPassPipelinePlan::getPerformancePassPipeline(const SILOptions &Options) {
   //
   // FIXME: When *not* emitting a .swiftmodule, skip the high-level function
   // pipeline to save compile time.
+  //
+  // NOTE: Ownership is now stripped within this function!
   addHighLevelFunctionPipeline(P);
 
   addHighLevelModulePipeline(P);

lib/Serialization/Deserialization.cpp

@@ -4371,16 +4371,26 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
       case decls_block::Derivative_DECL_ATTR: {
         bool isImplicit;
         uint64_t origNameId;
+        bool hasAccessorKind;
+        uint64_t rawAccessorKind;
         DeclID origDeclId;
         uint64_t rawDerivativeKind;
         ArrayRef<uint64_t> parameters;
 
         serialization::decls_block::DerivativeDeclAttrLayout::readRecord(
-            scratch, isImplicit, origNameId, origDeclId, rawDerivativeKind,
-            parameters);
+            scratch, isImplicit, origNameId, hasAccessorKind, rawAccessorKind,
+            origDeclId, rawDerivativeKind, parameters);
+
+        Optional<AccessorKind> accessorKind = None;
+        if (hasAccessorKind) {
+          auto maybeAccessorKind = getActualAccessorKind(rawAccessorKind);
+          if (!maybeAccessorKind)
+            MF.fatal();
+          accessorKind = *maybeAccessorKind;
+        }
 
         DeclNameRefWithLoc origName{DeclNameRef(MF.getDeclBaseName(origNameId)),
-                                    DeclNameLoc(), None};
+                                    DeclNameLoc(), accessorKind};
         auto derivativeKind =
             getActualAutoDiffDerivativeFunctionKind(rawDerivativeKind);
         if (!derivativeKind)

lib/Serialization/ModuleFormat.h

@@ -55,7 +55,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 563; // unchecked_value_cast
+const uint16_t SWIFTMODULE_VERSION_MINOR = 564; // `@derivative` attribute accessor kind
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1848,6 +1848,8 @@ namespace decls_block {
     Derivative_DECL_ATTR,
     BCFixed<1>, // Implicit flag.
     IdentifierIDField, // Original name.
+    BCFixed<1>, // Has original accessor kind?
+    AccessorKindField, // Original accessor kind.
     DeclIDField, // Original function declaration.
     AutoDiffDerivativeFunctionKindField, // Derivative function kind.
     BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.

lib/Serialization/Serialization.cpp

@@ -2431,19 +2431,25 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       assert(attr->getOriginalFunction(ctx) &&
              "`@derivative` attribute should have original declaration set "
              "during construction or parsing");
-      auto origName = attr->getOriginalFunctionName().Name.getBaseName();
+      auto origDeclNameRef = attr->getOriginalFunctionName();
+      auto origName = origDeclNameRef.Name.getBaseName();
       IdentifierID origNameId = S.addDeclBaseNameRef(origName);
       DeclID origDeclID = S.addDeclRef(attr->getOriginalFunction(ctx));
       auto derivativeKind =
           getRawStableAutoDiffDerivativeFunctionKind(attr->getDerivativeKind());
+      uint8_t rawAccessorKind = 0;
+      auto origAccessorKind = origDeclNameRef.AccessorKind;
+      if (origAccessorKind)
+        rawAccessorKind = uint8_t(getStableAccessorKind(*origAccessorKind));
       auto *parameterIndices = attr->getParameterIndices();
       assert(parameterIndices && "Parameter indices must be resolved");
       SmallVector<bool, 4> paramIndicesVector;
       for (unsigned i : range(parameterIndices->getCapacity()))
         paramIndicesVector.push_back(parameterIndices->contains(i));
       DerivativeDeclAttrLayout::emitRecord(
           S.Out, S.ScratchRecord, abbrCode, attr->isImplicit(), origNameId,
-          origDeclID, derivativeKind, paramIndicesVector);
+          origAccessorKind.hasValue(), rawAccessorKind, origDeclID,
+          derivativeKind, paramIndicesVector);
       return;
     }
 

stdlib/public/Differentiation/ArrayDifferentiation.swift

@@ -42,6 +42,14 @@ where Element: Differentiable {
     return (base, { $0 })
   }
 
+  @usableFromInline
+  @derivative(of: base)
+  func _jvpBase() -> (
+    value: [Element], differential: (Array<Element>.TangentVector) -> TangentVector
+  ) {
+    return (base, { $0 })
+  }
+
   /// Creates a differentiable view of the given array.
   public init(_ base: [Element]) { self._base = base }
 
@@ -53,6 +61,14 @@ where Element: Differentiable {
     return (Array.DifferentiableView(base), { $0 })
   }
 
+  @usableFromInline
+  @derivative(of: init(_:))
+  static func _jvpInit(_ base: [Element]) -> (
+    value: Array.DifferentiableView, differential: (TangentVector) -> TangentVector
+  ) {
+    return (Array.DifferentiableView(base), { $0 })
+  }
+
   public typealias TangentVector =
     Array<Element.TangentVector>.DifferentiableView
 
@@ -191,6 +207,17 @@ extension Array where Element: Differentiable {
     return (self[index], pullback)
   }
 
+  @usableFromInline
+  @derivative(of: subscript)
+  func _jvpSubscript(index: Int) -> (
+    value: Element, differential: (TangentVector) -> Element.TangentVector
+  ) {
+    func differential(_ v: TangentVector) -> Element.TangentVector {
+      return v[index]
+    }
+    return (self[index], differential)
+  }
+
   @usableFromInline
   @derivative(of: +)
   static func _vjpConcatenate(_ lhs: Self, _ rhs: Self) -> (
@@ -210,8 +237,26 @@ extension Array where Element: Differentiable {
     }
     return (lhs + rhs, pullback)
   }
+
+  @usableFromInline
+  @derivative(of: +)
+  static func _jvpConcatenate(_ lhs: Self, _ rhs: Self) -> (
+    value: Self,
+    differential: (TangentVector, TangentVector) -> TangentVector
+  ) {
+    func differential(_ l: TangentVector, _ r: TangentVector) -> TangentVector {
+      precondition(
+        l.base.count == lhs.count && r.base.count == rhs.count, """
+          Tangent vectors with invalid count; expected to equal the \
+          operand counts \(lhs.count) and \(rhs.count)
+          """)
+      return .init(l.base + r.base)
+    }
+    return (lhs + rhs, differential)
+  }
 }
 
+
 extension Array where Element: Differentiable {
   @usableFromInline
   @derivative(of: append)
@@ -277,6 +322,17 @@ extension Array where Element: Differentiable {
       }
     )
   }
+
+  @usableFromInline
+  @derivative(of: init(repeating:count:))
+  static func _jvpInit(repeating repeatedValue: Element, count: Int) -> (
+    value: Self, differential: (Element.TangentVector) -> TangentVector
+  ) {
+    (
+      value: Self(repeating: repeatedValue, count: count),
+      differential: { v in TangentVector(.init(repeating: v, count: count)) }
+    )
+  }
 }
 
 //===----------------------------------------------------------------------===//
@@ -312,6 +368,27 @@ extension Array where Element: Differentiable {
     }
     return (value: values, pullback: pullback)
   }
+
+  @inlinable
+  @derivative(of: differentiableMap)
+  internal func _jvpDifferentiableMap<Result: Differentiable>(
+    _ body: @differentiable (Element) -> Result
+  ) -> (
+    value: [Result],
+    differential: (Array.TangentVector) -> Array<Result>.TangentVector
+  ) {
+    var values: [Result] = []
+    var differentials: [(Element.TangentVector) -> Result.TangentVector] = []
+    for x in self {
+      let (y, df) = valueWithDifferential(at: x, in: body)
+      values.append(y)
+      differentials.append(df)
+    }
+    func differential(_ tans: Array.TangentVector) -> Array<Result>.TangentVector {
+      .init(zip(tans.base, differentials).map { tan, df in df(tan) })
+    }
+    return (value: values, differential: differential)
+  }
 }
 
 extension Array where Element: Differentiable {
@@ -361,4 +438,33 @@ extension Array where Element: Differentiable {
       }
     )
   }
+
+  @inlinable
+  @derivative(of: differentiableReduce, wrt: (self, initialResult))
+  func _jvpDifferentiableReduce<Result: Differentiable>(
+    _ initialResult: Result,
+    _ nextPartialResult: @differentiable (Result, Element) -> Result
+  ) -> (value: Result,
+        differential: (Array.TangentVector, Result.TangentVector)
+          -> Result.TangentVector) {
+    var differentials:
+      [(Result.TangentVector, Element.TangentVector) -> Result.TangentVector]
+        = []
+    let count = self.count
+    differentials.reserveCapacity(count)
+    var result = initialResult
+    for element in self {
+      let (y, df) =
+        valueWithDifferential(at: result, element, in: nextPartialResult)
+      result = y
+      differentials.append(df)
+    }
+    return (value: result, differential: { dSelf, dInitial in
+      var dResult = dInitial
+      for (dElement, df) in zip(dSelf.base, differentials) {
+        dResult = df(dResult, dElement)
+      }
+      return dResult
+    })
+  }
 }