[AutoDiff upstream] Add VJPEmitter.

`VJPEmitter` is a cloner that emits VJP functions. It implements reverse-mode
automatic differentiation, along with `PullbackEmitter`.

`VJPEmitter` clones an original function, replacing function applications with
VJP function applications. In VJP functions, each basic block takes a pullback
struct (containing callee pullbacks) and produces a predecessor enum: these data
structures are consumed by pullback functions.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -510,6 +510,25 @@ NOTE(autodiff_member_subset_indices_not_differentiable,none,
 NOTE(autodiff_cannot_param_subset_thunk_partially_applied_orig_fn,none,
      "cannot convert a direct method reference to a '@differentiable' "
      "function; use an explicit closure instead", ())
+NOTE(autodiff_cannot_differentiate_through_multiple_results,none,
+     "cannot differentiate through multiple results", ())
+// TODO(TF-1149): Remove this diagnostic.
+NOTE(autodiff_loadable_value_addressonly_tangent_unsupported,none,
+     "cannot yet differentiate value whose type %0 has a compile-time known "
+     "size, but whose 'TangentVector' contains stored properties of unknown "
+     "size; consider modifying %1 to use fewer generic parameters in stored "
+     "properties", (Type, Type))
+NOTE(autodiff_enums_unsupported,none,
+     "differentiating enum values is not yet supported", ())
+NOTE(autodiff_stored_property_no_corresponding_tangent,none,
+     "property cannot be differentiated because '%0.TangentVector' does not "
+     "have a member named '%1'", (StringRef, StringRef))
+NOTE(autodiff_coroutines_not_supported,none,
+     "differentiation of coroutine calls is not yet supported", ())
+NOTE(autodiff_cannot_differentiate_writes_to_global_variables,none,
+     "cannot differentiate writes to global variables", ())
+NOTE(autodiff_cannot_differentiate_writes_to_mutable_captures,none,
+     "cannot differentiate writes to mutable captures", ())
 
 ERROR(non_physical_addressof,none,
       "addressof only works with purely physical lvalues; "

include/swift/SIL/SILCloner.h

@@ -47,16 +47,16 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   TypeSubstitutionMap OpenedExistentialSubs;
   SILOpenedArchetypesTracker OpenedArchetypesTracker;
 
-private:
-  /// MARK: Private state hidden from CRTP extensions.
-
   // The old-to-new value map.
   llvm::DenseMap<SILValue, SILValue> ValueMap;
 
   /// The old-to-new block map. Some entries may be premapped with original
   /// blocks.
   llvm::DenseMap<SILBasicBlock*, SILBasicBlock*> BBMap;
 
+private:
+  /// MARK: Private state hidden from CRTP extensions.
+
   // The original blocks in DFS preorder. All blocks in this list are mapped.
   // After cloning, this represents the entire cloned CFG.
   //

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

@@ -0,0 +1,153 @@
+//===--- VJPEmitter.h - VJP Generation in differentiation -----*- C++ -*---===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a helper class for generating VJP functions for automatic
+// differentiation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_VJPEMITTER_H
+#define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_VJPEMITTER_H
+
+#include "swift/SIL/TypeSubstCloner.h"
+#include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
+#include "swift/SILOptimizer/Utils/Differentiation/DifferentiationInvoker.h"
+#include "swift/SILOptimizer/Utils/Differentiation/LinearMapInfo.h"
+#include "llvm/ADT/DenseMap.h"
+
+namespace swift {
+
+class SILDifferentiabilityWitness;
+class SILBasicBlock;
+class SILFunction;
+class SILInstruction;
+
+namespace autodiff {
+
+class ADContext;
+class PullbackEmitter;
+
+class VJPEmitter final
+    : public TypeSubstCloner<VJPEmitter, SILOptFunctionBuilder> {
+  friend class PullbackEmitter;
+
+private:
+  /// The global context.
+  ADContext &context;
+
+  /// The original function.
+  SILFunction *const original;
+
+  /// The differentiability witness.
+  SILDifferentiabilityWitness *const witness;
+
+  /// The VJP function.
+  SILFunction *const vjp;
+
+  /// The pullback function.
+  SILFunction *pullback;
+
+  /// The differentiation invoker.
+  DifferentiationInvoker invoker;
+
+  /// Info from activity analysis on the original function.
+  const DifferentiableActivityInfo &activityInfo;
+
+  /// The linear map info.
+  LinearMapInfo pullbackInfo;
+
+  /// Caches basic blocks whose phi arguments have been remapped (adding a
+  /// predecessor enum argument).
+  SmallPtrSet<SILBasicBlock *, 4> remappedBasicBlocks;
+
+  bool errorOccurred = false;
+
+  /// Mapping from original blocks to pullback values. Used to build pullback
+  /// struct instances.
+  llvm::DenseMap<SILBasicBlock *, SmallVector<SILValue, 8>> pullbackValues;
+
+  ASTContext &getASTContext() const { return vjp->getASTContext(); }
+  SILModule &getModule() const { return vjp->getModule(); }
+  const SILAutoDiffIndices getIndices() const {
+    return witness->getSILAutoDiffIndices();
+  }
+
+  static SubstitutionMap getSubstitutionMap(SILFunction *original,
+                                            SILFunction *vjp);
+
+  static const DifferentiableActivityInfo &
+  getActivityInfo(ADContext &context, SILFunction *original,
+                  SILAutoDiffIndices indices, SILFunction *vjp);
+
+public:
+  explicit VJPEmitter(ADContext &context, SILFunction *original,
+                      SILDifferentiabilityWitness *witness, SILFunction *vjp,
+                      DifferentiationInvoker invoker);
+
+  SILFunction *createEmptyPullback();
+
+  /// Run VJP generation. Returns true on error.
+  bool run();
+
+  void postProcess(SILInstruction *orig, SILInstruction *cloned);
+
+  /// Remap original basic blocks, adding predecessor enum arguments.
+  SILBasicBlock *remapBasicBlock(SILBasicBlock *bb);
+
+  /// General visitor for all instructions. If any error is emitted by previous
+  /// visits, bail out.
+  void visit(SILInstruction *inst);
+
+  void visitSILInstruction(SILInstruction *inst);
+
+private:
+  /// Get the lowered SIL type of the given AST type.
+  SILType getLoweredType(Type type);
+
+  /// Get the lowered SIL type of the given nominal type declaration.
+  SILType getNominalDeclLoweredType(NominalTypeDecl *nominal);
+
+  /// Build a pullback struct value for the original block corresponding to the
+  /// given terminator.
+  StructInst *buildPullbackValueStructValue(TermInst *termInst);
+
+  /// Build a predecessor enum instance using the given builder for the given
+  /// original predecessor/successor blocks and pullback struct value.
+  EnumInst *buildPredecessorEnumValue(SILBuilder &builder,
+                                      SILBasicBlock *predBB,
+                                      SILBasicBlock *succBB,
+                                      SILValue pbStructVal);
+
+public:
+  void visitReturnInst(ReturnInst *ri);
+
+  void visitBranchInst(BranchInst *bi);
+
+  void visitCondBranchInst(CondBranchInst *cbi);
+
+  void visitSwitchEnumInstBase(SwitchEnumInstBase *inst);
+
+  void visitSwitchEnumInst(SwitchEnumInst *sei);
+
+  void visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai);
+
+  // If an `apply` has active results or active inout arguments, replace it
+  // with an `apply` of its VJP.
+  void visitApplyInst(ApplyInst *ai);
+
+  void visitDifferentiableFunctionInst(DifferentiableFunctionInst *dfi);
+};
+
+} // end namespace autodiff
+} // end namespace swift
+
+#endif // SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_VJPEMITTER_H

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -40,6 +40,7 @@
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/Differentiation/ADContext.h"
 #include "swift/SILOptimizer/Utils/Differentiation/Thunk.h"
+#include "swift/SILOptimizer/Utils/Differentiation/VJPEmitter.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/BreadthFirstIterator.h"
@@ -944,9 +945,8 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     auto *vjp = createEmptyVJP(context, original, witness, serializeFunctions);
     witness->setVJP(vjp);
     context.recordGeneratedFunction(vjp);
-    // TODO(TF-1211): Upstream and use `VJPEmitter`. Fatal error with a nice
-    // message for now.
-    emitFatalError(context, vjp, "_fatalErrorVJPNotGenerated");
+    VJPEmitter emitter(context, original, witness, vjp, invoker);
+    return emitter.run();
   }
   return false;
 }

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -92,6 +92,9 @@ static void addMandatoryOptPipeline(SILPassPipelinePlan &P) {
   P.addAllocBoxToStack();
   P.addNoReturnFolding();
   addDefiniteInitialization(P);
+
+  // Automatic differentiation: canonicalize all differentiability witnesses
+  // and `differentiable_function` instructions.
   P.addDifferentiation();
 
   // Only run semantic arc opts if we are optimizing and if mandatory semantic

lib/SILOptimizer/Utils/Differentiation/CMakeLists.txt

@@ -4,4 +4,5 @@ silopt_register_sources(
   DifferentiationInvoker.cpp
   LinearMapInfo.cpp
   Thunk.cpp
+  VJPEmitter.cpp
 )