---

yaml
---
r: 262079
b: refs/heads/tensorflow
c: 2ab4b8a
h: refs/heads/master
i:
  262077: d30379e
  262075: 421798a
  262071: afc2ce4
  262063: fead07f
  262047: 2bceee3
  262015: 1263945

Author: rxwei

[refs]

@@ -818,7 +818,7 @@ refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-04-25-a: 22f738a831d43aff2b9c9773bcb65
 refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-05-08-a: 7d98cc16689baba5c8a3b90a9329bdcc1a12b4e9
 refs/heads/cherr42: a566ad54b073c2c56ac0a705d0a5bed9743135a5
 "refs/heads/codable_test_comment_fix": fc8f6824f7f347e1e8db55bff62db385c5728b5a
-refs/heads/tensorflow: 1cd29f7f6b9d69819306c0bb92428e0e144084fe
+refs/heads/tensorflow: 2ab4b8a8e7842ac52ea835932d132f82b7990753
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-11-a: 8126fd7a652e2f70ad6d76505239e34fb2ef3e1a
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-12-a: b3fd3dd84df6717f2e2e9df58c6d7e99fed57086
 refs/tags/swift-4.1-DEVELOPMENT-SNAPSHOT-2018-05-13-a: 71135119579039dc321c5f65d870050fe36efda2

branches/tensorflow/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp

@@ -265,6 +265,11 @@ struct DifferentiationInvoker {
     // a Swift AST node.
     GradientInst,
 
+    // No known invoker. This is the case when the differentiation is requested
+    // from SIL source via a `autodiff_function` instruction **without** being
+    // linked to a Swift AST node.
+    AutoDiffFunctionInst,
+
     // Invoked by the indirect application of differentiation. This case has an
     // associated differentiation task reference.
     IndirectDifferentiation,
@@ -287,10 +292,14 @@ struct DifferentiationInvoker {
 private:
   Kind kind;
   union Value {
-    /// The instruction associated with the `SILSource` case.
+    /// The instruction associated with the `GradientInst` case.
     GradientInst *gradientInst;
     Value(GradientInst *inst) : gradientInst(inst) {}
 
+    /// The instruction associated with the `AutoDiffFunctionInst` case.
+    AutoDiffFunctionInst *adFuncInst;
+    Value(AutoDiffFunctionInst *inst) : adFuncInst(inst) {}
+
     /// The parent differentiation task associated with the
     /// `IndirectDifferentiation` case.
     std::pair<ApplyInst *, DifferentiationTask *> indirectDifferentiation;
@@ -322,6 +331,8 @@ struct DifferentiationInvoker {
 public:
   DifferentiationInvoker(GradientInst *inst)
       : kind(Kind::GradientInst), value(inst) {}
+  DifferentiationInvoker(AutoDiffFunctionInst *inst)
+      : kind(Kind::AutoDiffFunctionInst), value(inst) {}
   DifferentiationInvoker(ApplyInst *applyInst, DifferentiationTask *task)
       : kind(Kind::IndirectDifferentiation), value(applyInst, task) {}
   DifferentiationInvoker(ReverseAutoDiffExpr *expr)
@@ -338,6 +349,11 @@ struct DifferentiationInvoker {
     return value.gradientInst;
   }
 
+  AutoDiffFunctionInst *getAutoDiffFunctionInst() const {
+    assert(kind == Kind::AutoDiffFunctionInst);
+    return value.adFuncInst;
+  }
+
   std::pair<ApplyInst *, DifferentiationTask *>
   getIndirectDifferentiation() const {
     assert(kind == Kind::IndirectDifferentiation);
@@ -812,6 +828,9 @@ void DifferentiationInvoker::print(llvm::raw_ostream &os) const {
   case Kind::GradientInst:
     os << "gradient_inst=(" << *getGradientInst() << ")";
     break;
+  case Kind::AutoDiffFunctionInst:
+    os << "autodiff_function_inst=(" << *getAutoDiffFunctionInst() << ")";
+    break;
   case Kind::IndirectDifferentiation: {
     auto indDiff = getIndirectDifferentiation();
     os << "indirect_differentiation=(apply_inst=(" << *indDiff.first
@@ -1177,6 +1196,7 @@ void ADContext::emitNondifferentiabilityError(SILInstruction *inst,
   // associated with any source location, we emit a diagnostic at the
   // instruction source location.
   case DifferentiationInvoker::Kind::GradientInst:
+  case DifferentiationInvoker::Kind::AutoDiffFunctionInst:
   case DifferentiationInvoker::Kind::SILDifferentiableAttribute:
     diagnose(opLoc,
              diag.getValueOr(diag::autodiff_expression_is_not_differentiable));
@@ -1293,7 +1313,7 @@ class DifferentiableActivityInfo;
 /// does not matter for the final result.
 ///
 /// Reference:
-/// Laurent Hascoët. Automatic Differentiation by Program Transformation. 2017.
+/// Laurent Hascoët. Automatic Differentiation by Program Transformation. 2007.
 class DifferentiableActivityAnalysis
     : public FunctionAnalysisBase<DifferentiableActivityInfo> {
 private:
@@ -4673,6 +4693,14 @@ static ReverseAutoDiffExpr *findDifferentialOperator(GradientInst *inst) {
   return inst->getLoc().getAsASTNode<ReverseAutoDiffExpr>();
 }
 
+/// Given an `autodiff_function` instruction, find the corresponding
+/// differential operator used in the AST. If no differential operator is found,
+/// return nullptr.
+static AutoDiffFunctionExpr *findDifferentialOperator(
+    AutoDiffFunctionInst *inst) {
+  return inst->getLoc().getAsASTNode<AutoDiffFunctionExpr>();
+}
+
 // Retrieve or create an empty gradient function based on a `gradient`
 // instruction and replace all users of the `gradient` instruction with the
 // gradient function. Returns the gradient function.
@@ -4953,6 +4981,9 @@ class Differentiation : public SILModuleTransform {
   /// pushing a differentiation task onto the global list. Returns true if any
   /// error occurred.
   bool processGradientInst(GradientInst *gi, ADContext &context);
+
+  bool processAutoDiffFunctionInst(AutoDiffFunctionInst *adfi,
+                                   ADContext &context);
 };
 } // end anonymous namespace
 
@@ -5003,6 +5034,60 @@ bool Differentiation::processGradientInst(GradientInst *gi,
   return false;
 }
 
+bool Differentiation::processAutoDiffFunctionInst(AutoDiffFunctionInst *adfi,
+                                                  ADContext &context) {
+  SILFunction *parent = adfi->getFunction();
+  auto origFnOperand = adfi->getOriginalFunction();
+  // If it traces back to a `function_ref`, differentiate that.
+  if (auto *originalFRI = findReferenceToVisibleFunction(origFnOperand)) {
+    auto *original = originalFRI->getReferencedFunction();
+    // TODO: Find syntax-level AD invoker from `adfi`.
+    auto *task = context.lookUpOrRegisterDifferentiationTask(
+        original, SILAutoDiffIndices(0, adfi->getParameterIndices()),
+        DifferentiationInvoker(adfi));
+    // Expand the `autodiff_function` instruction by adding the JVP and VJP
+    // functions.
+    SILBuilder builder(adfi);
+    auto loc = parent->getLocation();
+    auto *vjp = task->getVJP();
+    auto *vjpRef = builder.createFunctionRef(loc, vjp);
+    auto finalVJP = reapplyFunctionConversion(
+        vjpRef, originalFRI, origFnOperand, builder, loc);
+    SILValue finalJVP;
+    // TODO: Implement "forward-mode differentiation" to get JVP. Currently it's
+    // just an undef because we won't use it.
+    {
+      auto origFnTy = origFnOperand->getType().getAs<SILFunctionType>();
+      auto jvpType = origFnTy->getAutoDiffAssociatedFunctionType(
+          adfi->getParameterIndices(), /*resultIndex*/ 0,
+          /*differentiationOrder*/ 1, AutoDiffAssociatedFunctionKind::JVP,
+          *getModule(),
+          /*lookupConformance*/
+              LookUpConformanceInModule(getModule()->getSwiftModule()));
+      finalJVP = SILUndef::get(SILType::getPrimitiveObjectType(jvpType),
+                             getModule());
+    }
+    auto *newADFI = builder.createAutoDiffFunction(
+        loc, adfi->getParameterIndices(), adfi->getDifferentiationOrder(),
+        origFnOperand, {finalJVP, finalVJP});
+    adfi->replaceAllUsesWith(newADFI);
+    adfi->eraseFromParent();
+  }
+  // Differentiating opaque functions is not supported yet.
+  else {
+    // Find the original differential operator expression. Show an error at the
+    // operator, highlight the argument, and show a note at the definition site
+    // of the argument.
+    if (auto *expr = findDifferentialOperator(adfi))
+      context.diagnose(expr->getSubExpr()->getLoc(),
+                       diag::autodiff_function_not_differentiable)
+          .highlight(expr->getSubExpr()->getSourceRange());
+    return true;
+  }
+  PM->invalidateAnalysis(parent, SILAnalysis::InvalidationKind::FunctionBody);
+  return false;
+}
+
 /// AD pass entry.
 void Differentiation::run() {
   auto &module = *getModule();
@@ -5013,6 +5098,7 @@ void Differentiation::run() {
   SmallVector<std::pair<SILFunction *,
                         SILDifferentiableAttr *>, 8> diffAttrs;
   SmallVector<GradientInst *, 16> gradInsts;
+  SmallVector<AutoDiffFunctionInst *, 16> autodiffInsts;
   // Handle each `gradient` instruction and each `differentiable`
   // attribute in the module.
   for (SILFunction &f : module) {
@@ -5037,6 +5123,8 @@ void Differentiation::run() {
         // operator, push it to the work list.
         if (auto *gi = dyn_cast<GradientInst>(&i))
           gradInsts.push_back(gi);
+        else if (auto *adfi = dyn_cast<AutoDiffFunctionInst>(&i))
+          autodiffInsts.push_back(adfi);
       }
     }
   }
@@ -5081,9 +5169,12 @@ void Differentiation::run() {
   // turned into a differentiation task. But we don't back out just yet - primal
   // synthesis and adjoint synthesis for newly created differentiation tasks
   // should still run because they will diagnose more errors.
-  bool errorProcessingGradInsts = false;
+  bool errorProcessingAutoDiffInsts = false;
   for (auto *gi : gradInsts)
-    errorProcessingGradInsts |= processGradientInst(gi, context);
+    errorProcessingAutoDiffInsts |= processGradientInst(gi, context);
+
+  for (auto *adfi : autodiffInsts)
+    errorProcessingAutoDiffInsts |= processAutoDiffFunctionInst(adfi, context);
 
   // Run primal generation for newly created differentiation tasks. If any error
   // occurs, back out.
@@ -5099,7 +5190,7 @@ void Differentiation::run() {
 
   // If there was any error that occurred during `gradient` instruction
   // processing, back out.
-  if (errorProcessingGradInsts)
+  if (errorProcessingAutoDiffInsts)
     return;
 
   // Fill the body of each empty canonical gradient function.

branches/tensorflow/test/AutoDiff/autodiff_basic.sil

@@ -17,23 +17,30 @@ bb0(%0 : @trivial $Float):
 
 sil hidden @bar : $@convention(thin) (Float) -> (Float, Float) {
 bb0(%0 : @trivial $Float):
-  %1 = function_ref @foo : $@convention(thin) (Float) -> Float
-  %2 = gradient [source 0] [wrt 0] %1 : $@convention(thin) (Float) -> Float
-  %3 = apply %2(%0) : $@convention(thin) (Float) -> Float
-  %4 = gradient [source 0] [wrt 0] [preserving_result] %1 : $@convention(thin) (Float) -> Float
-  %5 = apply %4(%3) : $@convention(thin) (Float) -> (Float, Float)
-  return %5 : $(Float, Float)
+  %fref = function_ref @foo : $@convention(thin) (Float) -> Float
+
+  %grad = gradient [source 0] [wrt 0] %fref : $@convention(thin) (Float) -> Float
+  %grad_result = apply %grad(%0) : $@convention(thin) (Float) -> Float
+
+  %value_preserving_grad = gradient [source 0] [wrt 0] [preserving_result] %fref : $@convention(thin) (Float) -> Float
+  %value_and_grad = apply %value_preserving_grad(%0) : $@convention(thin) (Float) -> (Float, Float)
+
+  %adfunc = autodiff_function [wrt 0] [order 1] %fref : $@convention(thin) (Float) -> Float
+
+  return %value_and_grad : $(Float, Float)
 }
 
 // Here all `gradient` instructions have been replaced by `function_ref`s.
 
 // CHECK-LABEL: sil hidden @bar :
-// CHECK: bb0
-// CHECK:   %1 = function_ref @AD__foo__grad_src_0_wrt_0 : $@convention(thin) (Float) -> Float
-// CHECK:   %2 = apply %1(%0) : $@convention(thin) (Float) -> Float
-// CHECK:   %3 = function_ref @AD__foo__grad_src_0_wrt_0_p : $@convention(thin) (Float) -> (Float, Float)
-// CHECK:   %4 = apply %3(%2) : $@convention(thin) (Float) -> (Float, Float)
-// CHECK:   return %4 : $(Float, Float)
+// CHECK:   [[FREF:%.*]] = function_ref @foo : $@convention(thin) (Float) -> Float
+// CHECK:   [[GRAD_REF:%.*]] = function_ref @AD__foo__grad_src_0_wrt_0 : $@convention(thin) (Float) -> Float
+// CHECK:   [[GRAD_RESULT:%.*]] = apply [[GRAD_REF]](%0) : $@convention(thin) (Float) -> Float
+// CHECK:   [[VALUE_PRESERVING_GRAD_REF:%.*]] = function_ref @AD__foo__grad_src_0_wrt_0_p : $@convention(thin) (Float) -> (Float, Float)
+// CHECK:   [[VALUE_PRESERVING_GRAD_RESULT:%.*]] = apply [[VALUE_PRESERVING_GRAD_REF]](%0) : $@convention(thin) (Float) -> (Float, Float)
+// CHECK:   [[VJP_REF:%.*]] = function_ref @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK:   autodiff_function [wrt 0] [order 1] [[FREF]] : $@convention(thin) (Float) -> Float with {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[VJP_REF]] : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
+// CHECK:   return [[VALUE_PRESERVING_GRAD_RESULT]] : $(Float, Float)
 // CHECK: }
 
 // CHECK-LABEL:sil hidden @AD__foo__grad_src_0_wrt_0_s_p :

branches/tensorflow/test/AutoDiff/autodiff_function_inst_irgen.sil

@@ -16,7 +16,7 @@ bb0(%0 : @trivial $Float, %1 : @trivial $Float, %2 : @trivial $Float):
 }
 
 // The original function with an attribute that specifies the compiler-emitted pullback.
-sil hidden [differentiable source 0 wrt 0 primal @foo adjoint @foo_adj primitive] @foo : $@convention(thin) (Float) -> Float {
+sil hidden [differentiable source 0 wrt 0 primal @foo adjoint @foo_adj primitive vjp @foo_vjp] @foo : $@convention(thin) (Float) -> Float {
 bb0(%0 : @trivial $Float):
   return %0 : $Float
 }