Merge remote-tracking branch 'origin/tensorflow' into tensorflow-merge

Author: Eugene Burmako

include/swift/AST/ASTContext.h

@@ -110,8 +110,11 @@ namespace swift {
   class VarDecl;
   class UnifiedStatsReporter;
   class IndexSubset;
+  // SWIFT_ENABLE_TENSORFLOW
+  struct AutoDiffConfig;
   class VectorSpace;
   class DifferentiableAttr;
+  // SWIFT_ENABLE_TENSORFLOW END
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -711,6 +714,21 @@ class ASTContext final {
                        unsigned previousGeneration,
                        llvm::TinyPtrVector<AbstractFunctionDecl *> &methods);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Load derivative function configurations for the given
+  /// AbstractFunctionDecl.
+  ///
+  /// \param originalAFD The declaration whose derivative function
+  /// configurations should be loaded.
+  ///
+  /// \param previousGeneration The previous generation number. The AST already
+  /// contains derivative function configurations loaded from any generation up
+  /// to and including this one.
+  void loadDerivativeFunctionConfigurations(
+      AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+      llvm::SetVector<AutoDiffConfig> &results);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// Retrieve the Clang module loader for this ASTContext.
   ///
   /// If there is no Clang module loader, returns a null pointer.

include/swift/AST/Decl.h

@@ -5698,6 +5698,30 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
 private:
   ParameterList *Params;
 
+// SWIFT_ENABLE_TENSORFLOW
+private:
+  /// The generation at which we last loaded derivative function configurations.
+  unsigned DerivativeFunctionConfigGeneration = 0;
+  /// Prepare to traverse the list of derivative function configurations.
+  void prepareDerivativeFunctionConfigurations();
+
+  /// A uniqued list of derivative function configurations.
+  /// - `@differentiable` and `@derivative` attribute type-checking is
+  ///   responsible for populating derivative function configurations specified
+  ///   in the current module.
+  /// - Module loading is responsible for populating derivative function
+  ///   configurations from imported modules.
+  struct DerivativeFunctionConfigurationList;
+  DerivativeFunctionConfigurationList *DerivativeFunctionConfigs = nullptr;
+
+public:
+  /// Get all derivative function configurations.
+  ArrayRef<AutoDiffConfig> getDerivativeFunctionConfigurations();
+
+  /// Add the given derivative function configuration.
+  void addDerivativeFunctionConfiguration(AutoDiffConfig config);
+// SWIFT_ENABLE_TENSORFLOW END
+
 protected:
   // If a function has a body at all, we have either a parsed body AST node or
   // we have saved the end location of the unparsed body.

include/swift/AST/DiagnosticsSIL.def

@@ -467,6 +467,12 @@ NOTE(autodiff_expression_not_differentiable_note,none,
 NOTE(autodiff_external_nondifferentiable_function,none,
      "cannot differentiate functions that have not been marked "
      "'@differentiable' and that are defined in other files", ())
+NOTE(autodiff_private_derivative_from_fragile,none,
+     "differentiated functions in "
+     "%select{'@inlinable' functions|default arguments}0 must be marked "
+     "'@differentiable' or have a public '@derivative'"
+     "%select{|; this is not possible with a closure, make a top-level "
+     "function instead}1", (unsigned, bool))
 NOTE(autodiff_nondifferentiable_argument,none,
      "cannot differentiate through a non-differentiable argument; do you want "
      "to use 'withoutDerivative(at:)'?", ())

include/swift/AST/ModuleLoader.h

@@ -34,6 +34,9 @@ class DependencyCollector;
 
 namespace swift {
 
+// SWIFT_ENABLE_TENSORFLOW
+struct AutoDiffConfig;
+// SWIFT_ENABLE_TENSORFLOW END
 class AbstractFunctionDecl;
 class ClangImporterOptions;
 class ClassDecl;
@@ -151,6 +154,25 @@ class ModuleLoader {
                  unsigned previousGeneration,
                  llvm::TinyPtrVector<AbstractFunctionDecl *> &methods) = 0;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Load derivative function configurations for the given
+  /// AbstractFunctionDecl.
+  ///
+  /// \param originalAFD The declaration whose derivative function
+  /// configurations should be loaded.
+  ///
+  /// \param previousGeneration The previous generation number. The AST already
+  /// contains derivative function configurations loaded from any generation up
+  /// to and including this one.
+  ///
+  /// \param results The result list of derivative function configurations.
+  /// This list will be extended with any methods found in subsequent
+  /// generations.
+  virtual void loadDerivativeFunctionConfigurations(
+      AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+      llvm::SetVector<AutoDiffConfig> &results) {};
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// Verify all modules loaded by this loader.
   virtual void verifyAllModules() { }
 };

include/swift/AST/Types.h

@@ -4303,12 +4303,82 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
 
   CanSILFunctionType getWithoutDifferentiability();
 
-  /// Returns the type of the derivative function.
+  /// Returns the type of the derivative function for the given parameter
+  /// indices, result index, derivative function kind, derivative function
+  /// generic signature (optional), and other auxiliary parameters.
+  ///
+  /// Preconditions:
+  /// - Parameters corresponding to parameter indices must conform to
+  ///   `Differentiable`.
+  /// - The result corresponding to the result index must conform to
+  ///   `Differentiable`.
+  ///
+  /// Typing rules, given:
+  /// - Original function type: $(T0, T1, ...) -> (R0, R1, ...)
+  ///
+  /// Terminology:
+  /// - The derivative of a `Differentiable`-conforming type has the
+  ///   `TangentVector` associated type. `TangentVector` is abbreviated as `Tan`
+  ///   below.
+  /// - "wrt" parameters refers to parameters indicated by the parameter
+  ///   indices.
+  /// - "wrt" result refers to the result indicated by the result index.
+  ///
+  /// JVP derivative type:
+  /// - Takes original parameters.
+  /// - Returns original results, followed by a differential function, which
+  ///   takes "wrt" parameter derivatives and returns a "wrt" result derivative.
+  ///
+  ///     $(T0, ...) -> (R0, ...,  (T0.Tan, T1.Tan, ...) -> R0.Tan)
+  ///                    ^~~~~~~    ^~~~~~~~~~~~~~~~~~~     ^~~~~~
+  ///          original results | derivatives wrt params | derivative wrt result
+  ///
+  /// VJP derivative type:
+  /// - Takes original parameters.
+  /// - Returns original results, followed by a pullback function, which
+  ///   takes a "wrt" result derivative and returns "wrt" parameter derivatives.
+  ///
+  ///     $(T0, ...) -> (R0, ...,       (R0.Tan)  ->     (T0.Tan, T1.Tan, ...))
+  ///                    ^~~~~~~         ^~~~~~            ^~~~~~~~~~~~~~~~~~~
+  ///          original results | derivative wrt result | derivatives wrt params
+  ///
+  /// The JVP/VJP generic signature is a "constrained" version of the given
+  /// `derivativeFunctionGenericSignature` if specified. Otherwise, it is a
+  /// "constrained" version of the original generic signature. A "constrained"
+  /// generic signature requires all "wrt" parameters to conform to
+  /// `Differentiable`; this is important for correctness.
+  ///
+  /// Other properties of the original function type are copied exactly:
+  /// `ExtInfo`, coroutine kind, callee convention, yields, optional error
+  /// result, witness method conformance, etc.
+  ///
+  /// Special cases:
+  /// - Reabstraction thunks have special derivative type calculation. The
+  ///   original function-typed last parameter is transformed into a
+  ///   `@differentiable` function-typed parameter in the derivative type. This
+  ///   is necessary for the differentiation transform to support reabstraction
+  ///   thunk differentiation because the function argument is opaque and cannot
+  ///   be differentiated. Instead, the argument is made `@differentiable` and
+  ///   reabstraction thunk JVP/VJP callers are responsible for passing a
+  ///   `@differentiable` function.
+  ///   - TODO(TF-1036): Investigate more efficient reabstraction thunk
+  ///     derivative approaches. The last argument can simply be a
+  ///     corresponding derivative function, instead of a `@differentiable`
+  ///     function - this is more direct. It may be possible to implement
+  ///     reabstraction thunk derivatives using "reabstraction thunks for
+  ///     the original function's derivative", avoiding extra code generation.
+  ///
+  /// Caveats:
+  /// - We may support multiple result indices instead of a single result index
+  ///   eventually. At the SIL level, this enables differentiating wrt multiple
+  ///   function results. At the Swift level, this enables differentiating wrt
+  ///   multiple tuple elements for tuple-returning functions.
   CanSILFunctionType getAutoDiffDerivativeFunctionType(
       IndexSubset *parameterIndices, unsigned resultIndex,
       AutoDiffDerivativeFunctionKind kind, Lowering::TypeConverter &TC,
       LookupConformanceFn lookupConformance,
-      CanGenericSignature derivativeFunctionGenericSignature = nullptr);
+      CanGenericSignature derivativeFunctionGenericSignature = nullptr,
+      bool isReabstractionThunk = false);
 
   /// Returns the type of the transpose function.
   CanSILFunctionType getAutoDiffTransposeFunctionType(

include/swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h

@@ -131,7 +131,7 @@ class DifferentiableActivityInfo {
   SmallVector<SmallDenseSet<SILValue>, 4> usefulValueSets;
 
   /// The original function.
-  SILFunction &getFunction();
+  SILFunction &getFunction() const;
 
   /// Returns true if the given SILValue has a tangent space.
   bool hasTangentSpace(SILValue value) {
@@ -206,6 +206,14 @@ class DifferentiableActivityInfo {
   /// Returns the activity of the given value for the given `SILAutoDiffIndices`
   /// (parameter indices and result index).
   Activity getActivity(SILValue value, const SILAutoDiffIndices &indices) const;
+
+  /// Prints activity information for the `indices` of the given `value`.
+  void dump(SILValue value, const SILAutoDiffIndices &indices,
+            llvm::raw_ostream &s = llvm::dbgs()) const;
+
+  /// Prints activity information for the given `indices`.
+  void dump(SILAutoDiffIndices indices,
+            llvm::raw_ostream &s = llvm::dbgs()) const;
 };
 
 class DifferentiableActivityCollection {

include/swift/SILOptimizer/Utils/Differentiation/ADContext.h

@@ -136,7 +136,8 @@ class ADContext {
   }
 
   /// Adds the given `differentiable_function` instruction to the worklist.
-  void addDifferentiableFunctionInst(DifferentiableFunctionInst* dfi) {
+  void
+  addDifferentiableFunctionInstToWorklist(DifferentiableFunctionInst *dfi) {
     differentiableFunctionInsts.push_back(dfi);
   }