---

yaml
---
r: 262075
b: refs/heads/tensorflow
c: 932f7fa
h: refs/heads/master
i:
  262073: 3667a78
  262071: afc2ce4

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: b435eef3e693c6665ce53a287026d01f48411ca8
+refs/heads/tensorflow: 932f7fa6ec4db24d344ac26198d25432d94dbf46
 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/include/swift/AST/Builtins.def

@@ -376,10 +376,10 @@ BUILTIN_SIL_OPERATION(AllocWithTailElems, "allocWithTailElems", Special)
 BUILTIN_SIL_OPERATION(ProjectTailElems, "projectTailElems", Special)
 
 // SWIFT_ENABLE_TENSORFLOW
-/// autodifGetJVP has type <T: Differentiable, R: Differentiable>
+/// autodiffGetJVP has type <T: Differentiable, R: Differentiable>
 ///     ((T) -> R) -> (T) -> (R, (T.TangentVector) -> R.TangentVector).
 BUILTIN_SIL_OPERATION(AutoDiffGetJVP, "autodiffGetJVP", Special)
-/// autodifGetVJP has type <T: Differentiable, R: Differentiable>
+/// autodiffGetVJP has type <T: Differentiable, R: Differentiable>
 ///     ((T) -> R) -> (T) -> (R, (R.CotangentVector) -> T.CotangentVector).
 BUILTIN_SIL_OPERATION(AutoDiffGetVJP, "autodiffGetVJP", Special)
 

branches/tensorflow/include/swift/SIL/SILInstruction.h

@@ -7648,13 +7648,14 @@ class AutoDiffFunctionExtractInst :
                            SingleValueInstruction> {
 public:
   struct Extractee {
-    enum innerty : uint8_t {
+    enum innerty : unsigned {
       Original = 0,
       JVP = 1,
       VJP = 2
     } rawValue;
     Extractee() = default;
     Extractee(innerty rawValue) : rawValue(rawValue) {}
+    Extractee(unsigned rawValue) : Extractee((innerty)rawValue) {}
     explicit Extractee(StringRef name);
     operator innerty() const { return rawValue; }
   };

branches/tensorflow/include/swift/Serialization/ModuleFormat.h

@@ -52,7 +52,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 = 457; // Last change: add jvp and vjp to sil differentiable attribute
+const uint16_t SWIFTMODULE_VERSION_MINOR = 458; // Last change: serialize autodiff_function and autodiff_function_extract
 
 using DeclIDField = BCFixed<31>;
 

branches/tensorflow/lib/SIL/OwnershipUtils.cpp

@@ -45,6 +45,8 @@ bool swift::isOwnershipForwardingValueKind(SILNodeKind kind) {
   case SILNodeKind::DestructureTupleInst:
   // SWIFT_ENABLE_TENSORFLOW
   case SILNodeKind::GradientInst:
+  case SILNodeKind::AutoDiffFunctionInst:
+  case SILNodeKind::AutoDiffFunctionExtractInst:
     return true;
   default:
     return false;

branches/tensorflow/lib/Serialization/DeserializeSIL.cpp

@@ -972,7 +972,9 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn, SILBasicBlock *BB,
   Builder.setInsertionPoint(BB);
   Builder.setCurrentDebugScope(Fn->getDebugScope());
   unsigned RawOpCode = 0, TyCategory = 0, TyCategory2 = 0, TyCategory3 = 0,
-           Attr = 0, NumSubs = 0, NumConformances = 0, IsNonThrowingApply = 0;
+           // SWIFT_ENABLE_TENSORFLOW
+           Attr = 0, Attr2 = 0, NumSubs = 0, NumConformances = 0,
+           IsNonThrowingApply = 0;
   // SWIFT_ENABLE_TENSORFLOW
   unsigned NumArguments = 0;
   unsigned GradResultIndex = 0;
@@ -1087,6 +1089,18 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn, SILBasicBlock *BB,
                                       GradResultIndex, ListOfValues);
     RawOpCode = (unsigned)SILInstructionKind::GradientInst;
     break;
+  case SIL_INST_AUTODIFF_FUNCTION:
+    SILInstAutoDiffFunctionLayout::readRecord(scratch, /*order*/ Attr,
+                                              /*numParams*/ Attr2, NumArguments,
+                                              ListOfValues);
+    RawOpCode = (unsigned)SILInstructionKind::AutoDiffFunctionInst;
+    break;
+  case SIL_INST_AUTODIFF_FUNCTION_EXTRACT:
+    SILInstAutoDiffFunctionExtractLayout::readRecord(scratch, TyID, TyCategory,
+                                                     ValID, /*extractee*/ Attr,
+                                                     /*order*/ Attr2);
+    RawOpCode = (unsigned)SILInstructionKind::AutoDiffFunctionExtractInst;
+    break;
   case SIL_INST_NO_OPERAND:
     SILInstNoOperandLayout::readRecord(scratch, RawOpCode);
     break;
@@ -1488,10 +1502,32 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn, SILBasicBlock *BB,
     break;
   }
   case SILInstructionKind::AutoDiffFunctionInst: {
-    llvm_unreachable("FIXME: Unhandled");
+    auto numParamIndices = ListOfValues.size() - NumArguments * 3;
+    auto paramIndices = ListOfValues.take_front(numParamIndices);
+    auto numParams = Attr2;
+    llvm::SmallBitVector paramIndicesBitVec(numParams);
+    for (unsigned idx : paramIndices)
+      paramIndicesBitVec.set(idx);
+    SmallVector<SILValue, 4> operands;
+    for (auto i = numParamIndices; i < NumArguments * 3; i += 3) {
+      auto astTy = MF->getType(ListOfValues[i]);
+      auto silTy = getSILType(astTy, (SILValueCategory)ListOfValues[i+1]);
+      operands.push_back(getLocalValue(ListOfValues[i+2], silTy));
+    }
+    ResultVal = Builder.createAutoDiffFunction(Loc, paramIndicesBitVec,
+        /*differentiationOrder*/ Attr, operands[0],
+        ArrayRef<SILValue>(operands).drop_front());
+    break;
   }
   case SILInstructionKind::AutoDiffFunctionExtractInst: {
-    llvm_unreachable("FIXME: unhandled");
+    auto astTy = MF->getType(TyID);
+    auto silTy = getSILType(astTy, SILValueCategory::Object);
+    auto val = getLocalValue(ValID, silTy);
+    AutoDiffFunctionExtractee extractee(Attr);
+    auto order = Attr2;
+    ResultVal =
+        Builder.createAutoDiffFunctionExtract(Loc, extractee, order, val);
+    break;
   }
   case SILInstructionKind::GraphOperationInst: {
     // TODO(SR-8848): Deserialize attributes.

branches/tensorflow/lib/Serialization/SILFormat.h

@@ -175,6 +175,8 @@ namespace sil_block {
     SIL_REVERSE_DIFFERENTIABLE_ATTR,
     SIL_INST_GRAPH_OPERATION,
     SIL_INST_GRADIENT,
+    SIL_INST_AUTODIFF_FUNCTION,
+    SIL_INST_AUTODIFF_FUNCTION_EXTRACT,
 
     // We also share these layouts from the decls block. Their enumerators must
     // not overlap with ours.
@@ -293,7 +295,7 @@ namespace sil_block {
                      BCFixed<3>,  // side effect info.
                      BCVBR<8>,    // number of specialize attributes
                      // SWIFT_ENABLE_TENSORFLOW
-                     BCVBR<8>,    // number of reverse differentiable attributes
+                     BCVBR<8>,    // number of differentiable attributes
                      BCFixed<1>,  // has qualified ownership
                      BCFixed<1>,  // must be weakly referenced
                      TypeIDField, // SILFunctionType
@@ -429,6 +431,23 @@ namespace sil_block {
     BCArray<BCFixed<1>> // parameter indices
   >;
 
+  using SILInstAutoDiffFunctionLayout = BCRecordLayout<
+    SIL_INST_AUTODIFF_FUNCTION,
+    BCVBR<8>,             // differentiation order
+    BCVBR<8>,             // number of function parameters
+    BCVBR<8>,             // number of operands
+    BCArray<ValueIDField> // parameter indices and operands
+  >;
+
+  using SILInstAutoDiffFunctionExtractLayout = BCRecordLayout<
+    SIL_INST_AUTODIFF_FUNCTION_EXTRACT,
+    TypeIDField,
+    SILTypeCategoryField,
+    ValueIDField,
+    BCFixed<2>, // extractee
+    BCVBR<8>    // order
+  >;
+
   // SIL instructions with one type. (alloc_stack)
   using SILOneTypeLayout = BCRecordLayout<
     SIL_ONE_TYPE,

branches/tensorflow/lib/Serialization/SerializeSIL.cpp

@@ -963,10 +963,34 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
     break;
   }
   case SILInstructionKind::AutoDiffFunctionInst: {
-    llvm_unreachable("FIXME: handle this");
+    auto *adfi = cast<AutoDiffFunctionInst>(&SI);
+    SmallVector<ValueID, 4> trailingInfo;
+    auto &paramIndices = adfi->getParameterIndices();
+    for (unsigned idx : paramIndices.set_bits())
+      trailingInfo.push_back(idx);
+    for (auto &op : adfi->getAllOperands()) {
+      auto val = op.get();
+      trailingInfo.push_back(S.addTypeRef(val->getType().getASTType()));
+      trailingInfo.push_back((unsigned)val->getType().getCategory());
+      trailingInfo.push_back(addValueRef(val));
+    }
+    SILInstAutoDiffFunctionLayout::emitRecord(Out, ScratchRecord,
+        SILAbbrCodes[SILInstAutoDiffFunctionLayout::Code],
+        adfi->getDifferentiationOrder(), (unsigned)paramIndices.size(),
+        adfi->getNumOperands(), trailingInfo);
+    break;
   }
   case SILInstructionKind::AutoDiffFunctionExtractInst: {
-    llvm_unreachable("FIXME: handle this");
+    auto *adfei = cast<AutoDiffFunctionExtractInst>(&SI);
+    auto operandRef = addValueRef(adfei->getFunctionOperand());
+    auto operandType = adfei->getFunctionOperand()->getType();
+    auto operandTypeRef = S.addTypeRef(operandType.getASTType());
+    auto rawExtractee = (unsigned)adfei->getExtractee();
+    SILInstAutoDiffFunctionExtractLayout::emitRecord(Out, ScratchRecord,
+        SILAbbrCodes[SILInstAutoDiffFunctionExtractLayout::Code],
+        operandTypeRef, (unsigned)operandType.getCategory(), operandRef,
+        rawExtractee, adfei->getDifferentiationOrder());
+    break;
   }
   case SILInstructionKind::GraphOperationInst: {
     // TODO(SR-8848): Serialize attributes.
@@ -2461,6 +2485,8 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
   registerSILAbbr<SILDifferentiableAttrLayout>();
   registerSILAbbr<SILInstGraphOperationLayout>();
   registerSILAbbr<SILInstGradientLayout>();
+  registerSILAbbr<SILInstAutoDiffFunctionLayout>();
+  registerSILAbbr<SILInstAutoDiffFunctionExtractLayout>();
 
   // Register the abbreviation codes so these layouts can exist in both
   // decl blocks and sil blocks.

branches/tensorflow/stdlib/public/core/AutoDiff.swift

@@ -94,20 +94,65 @@ public extension Differentiable where TangentVector == CotangentVector {
 // Differential Operators
 //===----------------------------------------------------------------------===//
 
-@_transparent @_effects(readnone)
-public func gradient<T : Differentiable, R : FloatingPoint>(
-  of f: /*@autodiff*/ @escaping (T) -> R
-) -> (T) -> T {
-  return #gradient(f)
+// FIXME(rxwei): Fix SR-9458.
+#if false
+
+@inlinable
+public func valueWithDifferential<T, R>(
+  at x: T, in f: @escaping @autodiff (T) -> R
+) -> (value: R, differential: (T.TangentVector) -> R.TangentVector)
+  where T : Differentiable, R : Differentiable {
+  return Builtin.autodiffGetJVP(f)(x)
 }
 
-@_transparent @_effects(readnone)
-public func gradient<T : Differentiable, R : FloatingPoint>(
-  at x: T, in f: /*@autodiff*/ @escaping (T) -> R
-) -> T {
-  return #gradient(f)(x)
+@inlinable
+public func valueWithPullback<T, R>(
+  at x: T, in f: @escaping @autodiff (T) -> R
+) -> (value: R, pullback: (R.CotangentVector) -> T.CotangentVector)
+  where T : Differentiable, R : Differentiable {
+  return Builtin.autodiffGetVJP(f)(x)
 }
 
+@inlinable
+public func derivative<T, R>(
+  at x: T, in f: @escaping @autodiff (T) -> R
+) -> R.TangentVector
+  where T : BinaryFloatingPoint & Differentiable, R : Differentiable,
+        T.TangentVector == T {
+  let (y, differential) = valueWithDifferential(at: x, in: f)
+  return differential(1)
+}
+
+@inlinable
+public func derivative<T, R>(
+  of f: @escaping @autodiff (T) -> R
+) -> (T) -> R.TangentVector
+  where T : BinaryFloatingPoint & Differentiable, R : Differentiable,
+        T.TangentVector == T {
+  return { x in derivative(at: x, in: f) }
+}
+
+@inlinable
+public func gradient<T, R>(
+  at x: T, in f: @escaping @autodiff (T) -> R
+) -> T.CotangentVector
+  where T : Differentiable, R : BinaryFloatingPoint & Differentiable,
+        R.CotangentVector == R {
+  let (y, pullback) = valueWithPullback(at: x, in: f)
+  return pullback(1)
+}
+
+@inlinable
+public func gradient<T, R>(
+  of f: @escaping @autodiff (T) -> R
+) -> (T) -> T.CotangentVector
+  where T : Differentiable, R : BinaryFloatingPoint & Differentiable,
+        R.CotangentVector == R {
+  return { x in gradient(at: x, in: f) }
+}
+
+#endif
+
 //===----------------------------------------------------------------------===//
 // Builtins
 //===----------------------------------------------------------------------===//

branches/tensorflow/test/AutoDiff/autodiff_function_inst.sil

@@ -1,5 +1,10 @@
 // RUN: %target-sil-opt %s | %FileCheck %s
 
+// RUN: %empty-directory(%t)
+// RUN: %target-sil-opt %s -emit-sib -o %t/tmp.sib -module-name autodiff_function
+// RUN: %target-sil-opt %t/tmp.sib -o %t/tmp.2.sib -module-name autodiff_function
+// RUN: %target-sil-opt %t/tmp.2.sib -module-name autodiff_function | %FileCheck %s
+
 sil_stage raw
 
 import Swift