[AutoDiff] Support differentiation of loops.

dan-zheng · dan-zheng · commit 16dc9a2e9e9f · 2019-06-18T00:09:12.000-07:00
- Change predecessor enum generation to support loops.
  - Generate indirect predecessor enums for BBs in loops.
  - Handle boxed payloads of indirect enums.
- Traverse basic blocks in post-order post-dominance order.
  - This is necessary for computational correctness.
- Add loop tests (`for-in`, `while`, nested).
diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp
diff --git a/lib/SILOptimizer/Mandatory/Differentiation.h b/lib/SILOptimizer/Mandatory/Differentiation.h
@@ -0,0 +1,93 @@
+//===--- Differentiation.h - SIL Automatic Differentiation ----*- C++ -*---===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 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
+//
+//===----------------------------------------------------------------------===//
+//
+// SWIFT_ENABLE_TENSORFLOW
+//
+// Reverse-mode automatic differentiation utilities.
+//
+// NOTE: Although the AD feature is developed as part of the Swift for
+// TensorFlow project, it is completely independent from TensorFlow support.
+//
+// TODO: Move definitions here from Differentiation.cpp.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H
+#define SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H
+
+#include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
+#include "swift/SILOptimizer/Utils/Local.h"
+
+namespace swift {
+
+using llvm::DenseMap;
+using llvm::SmallDenseMap;
+using llvm::SmallDenseSet;
+using llvm::SmallMapVector;
+using llvm::SmallSet;
+
+/// Helper class for visiting basic blocks in post-order post-dominance order,
+/// based on a worklist algorithm.
+class PostOrderPostDominanceOrder {
+  SmallVector<DominanceInfoNode *, 16> buffer;
+  PostOrderFunctionInfo *postOrderInfo;
+  size_t srcIdx = 0;
+
+public:
+  /// Constructor.
+  /// \p root The root of the post-dominator tree.
+  /// \p postOrderInfo The post-order info of the function.
+  /// \p capacity Should be the number of basic blocks in the dominator tree to
+  ///             reduce memory allocation.
+  PostOrderPostDominanceOrder(DominanceInfoNode *root,
+                              PostOrderFunctionInfo *postOrderInfo,
+                              int capacity = 0)
+      : postOrderInfo(postOrderInfo) {
+    buffer.reserve(capacity);
+    buffer.push_back(root);
+  }
+
+  /// Get the next block from the worklist.
+  DominanceInfoNode *getNext() {
+    if (srcIdx == buffer.size())
+      return nullptr;
+    return buffer[srcIdx++];
+  }
+
+  /// Pushes the dominator children of a block onto the worklist in post-order.
+  void pushChildren(DominanceInfoNode *node) {
+    pushChildrenIf(node, [](SILBasicBlock *) { return true; });
+  }
+
+  /// Conditionally pushes the dominator children of a block onto the worklist
+  /// in post-order.
+  template <typename Pred>
+  void pushChildrenIf(DominanceInfoNode *node, Pred pred) {
+    SmallVector<DominanceInfoNode *, 4> children;
+    for (auto *child : *node)
+      children.push_back(child);
+    llvm::sort(children.begin(), children.end(),
+               [&](DominanceInfoNode *n1, DominanceInfoNode *n2) {
+                 return postOrderInfo->getPONumber(n1->getBlock()) <
+                        postOrderInfo->getPONumber(n2->getBlock());
+               });
+    for (auto *child : children) {
+      SILBasicBlock *childBB = child->getBlock();
+      if (pred(childBB))
+        buffer.push_back(child);
+    }
+  }
+};
+
+} // end namespace swift
+
+#endif // SWIFT_SILOPTIMIZER_MANDATORY_DIFFERENTIATION_H
diff --git a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
@@ -181,6 +181,16 @@ extension Tracked where T : Differentiable & SignedNumeric, T == T.Magnitude,
   }
 }
 
+extension Tracked where T : Differentiable & FloatingPoint,
+                        T == T.AllDifferentiableVariables, T == T.TangentVector {
+  @usableFromInline
+  @differentiating(/)
+  internal static func _vjpDivide(lhs: Self, rhs: Self)
+      -> (value: Self, pullback: (Self) -> (Self, Self)) {
+    return (lhs / rhs, { v in (v / rhs, -lhs / (rhs * rhs) * v) })
+  }
+}
+
 // Differential operators for `Tracked<Float>`.
 public extension Differentiable {
   @inlinable
diff --git a/test/AutoDiff/control_flow.swift b/test/AutoDiff/control_flow.swift
@@ -519,4 +519,46 @@ ControlFlowTests.test("Enums") {
   }
 }
 
+ControlFlowTests.test("Loops") {
+  func for_loop(_ x: Float) -> Float {
+    var result = x
+    for _ in 1..<3 {
+      result = result * x
+    }
+    return result
+  }
+  expectEqual((8, 12), valueWithGradient(at: 2, in: for_loop))
+  expectEqual((27, 27), valueWithGradient(at: 3, in: for_loop))
+
+  func while_loop(_ x: Float) -> Float {
+    var result = x
+    var i = 1
+    while i < 3 {
+      result = result * x
+      i += 1
+    }
+    return result
+  }
+  expectEqual((8, 12), valueWithGradient(at: 2, in: while_loop))
+  expectEqual((27, 27), valueWithGradient(at: 3, in: while_loop))
+
+  func nested_loop(_ x: Float) -> Float {
+    var outer = x
+    for _ in 1..<3 {
+      outer = outer * x
+
+      var inner = outer
+      var i = 1
+      while i < 3 {
+        inner = inner / x
+        i += 1
+      }
+      outer = inner
+    }
+    return outer
+  }
+  expectEqual((0.5, -0.25), valueWithGradient(at: 2, in: nested_loop))
+  expectEqual((0.25, -0.0625), valueWithGradient(at: 4, in: nested_loop))
+}
+
 runAllTests()
diff --git a/test/AutoDiff/control_flow_diagnostics.swift b/test/AutoDiff/control_flow_diagnostics.swift
@@ -33,6 +33,57 @@ func enum_nonactive2(_ e: Enum, _ x: Float) -> Float {
   }
 }
 
+// Test loops.
+
+@differentiable
+func for_loop(_ x: Float) -> Float {
+  var result: Float = x
+  for _ in 0..<3 {
+    result = result * x
+  }
+  return result
+}
+
+@differentiable
+func while_loop(_ x: Float) -> Float {
+  var result = x
+  var i = 1
+  while i < 3 {
+    result = result * x
+    i += 1
+  }
+  return result
+}
+
+@differentiable
+func nested_loop(_ x: Float) -> Float {
+  var outer = x
+  for _ in 1..<3 {
+    outer = outer * x
+
+    var inner = outer
+    var i = 1
+    while i < 3 {
+      inner = inner / x
+      i += 1
+    }
+    outer = inner
+  }
+  return outer
+}
+
+// Test `try_apply`.
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable
+// expected-note @+1 {{when differentiating this function definition}}
+func withoutDerivative<T : Differentiable, R: Differentiable>(
+  at x: T, in body: (T) throws -> R
+) rethrows -> R {
+  // expected-note @+1 {{differentiating control flow is not yet supported}}
+  try body(x)
+}
+
 // Test unsupported differentiation of active enum values.
 
 // expected-error @+1 {{function is not differentiable}}
@@ -90,17 +141,3 @@ enum Tree : Differentiable & AdditiveArithmetic {
     }
   }
 }
-
-// Test loops.
-
-// expected-error @+1 {{function is not differentiable}}
-@differentiable
-// expected-note @+1 {{when differentiating this function definition}}
-func loop(_ x: Float) -> Float {
-  var result: Float = 1
-  // expected-note @+1 {{differentiating loops is not yet supported}}
-  for _ in 0..<3 {
-    result += x
-  }
-  return x
-}
diff --git a/test/AutoDiff/control_flow_sil.swift b/test/AutoDiff/control_flow_sil.swift
@@ -16,31 +16,31 @@ func cond(_ x: Float) -> Float {
   return x - x
 }
 
-// CHECK-DATA-STRUCTURES: enum _AD__cond_bb0__Pred__src_0_wrt_0 {
-// CHECK-DATA-STRUCTURES: }
 // CHECK-DATA-STRUCTURES: struct _AD__cond_bb0__PB__src_0_wrt_0 {
 // CHECK-DATA-STRUCTURES: }
-// CHECK-DATA-STRUCTURES: enum _AD__cond_bb1__Pred__src_0_wrt_0 {
-// CHECK-DATA-STRUCTURES:   case bb0(_AD__cond_bb0__PB__src_0_wrt_0)
-// CHECK-DATA-STRUCTURES: }
 // CHECK-DATA-STRUCTURES: struct _AD__cond_bb1__PB__src_0_wrt_0 {
 // CHECK-DATA-STRUCTURES:   @_hasStorage var predecessor: _AD__cond_bb1__Pred__src_0_wrt_0 { get set }
 // CHECK-DATA-STRUCTURES:   @_hasStorage var pullback_0: (Float) -> (Float, Float) { get set }
 // CHECK-DATA-STRUCTURES: }
-// CHECK-DATA-STRUCTURES: enum _AD__cond_bb2__Pred__src_0_wrt_0 {
-// CHECK-DATA-STRUCTURES:   case bb0(_AD__cond_bb0__PB__src_0_wrt_0)
-// CHECK-DATA-STRUCTURES: }
 // CHECK-DATA-STRUCTURES: struct _AD__cond_bb2__PB__src_0_wrt_0 {
 // CHECK-DATA-STRUCTURES:   @_hasStorage var predecessor: _AD__cond_bb2__Pred__src_0_wrt_0 { get set }
 // CHECK-DATA-STRUCTURES:   @_hasStorage var pullback_1: (Float) -> (Float, Float) { get set }
 // CHECK-DATA-STRUCTURES: }
+// CHECK-DATA-STRUCTURES: struct _AD__cond_bb3__PB__src_0_wrt_0 {
+// CHECK-DATA-STRUCTURES:   @_hasStorage var predecessor: _AD__cond_bb3__Pred__src_0_wrt_0 { get set }
+// CHECK-DATA-STRUCTURES: }
+// CHECK-DATA-STRUCTURES: enum _AD__cond_bb0__Pred__src_0_wrt_0 {
+// CHECK-DATA-STRUCTURES: }
+// CHECK-DATA-STRUCTURES: enum _AD__cond_bb1__Pred__src_0_wrt_0 {
+// CHECK-DATA-STRUCTURES:   case bb0(_AD__cond_bb0__PB__src_0_wrt_0)
+// CHECK-DATA-STRUCTURES: }
+// CHECK-DATA-STRUCTURES: enum _AD__cond_bb2__Pred__src_0_wrt_0 {
+// CHECK-DATA-STRUCTURES:   case bb0(_AD__cond_bb0__PB__src_0_wrt_0)
+// CHECK-DATA-STRUCTURES: }
 // CHECK-DATA-STRUCTURES: enum _AD__cond_bb3__Pred__src_0_wrt_0 {
 // CHECK-DATA-STRUCTURES:   case bb2(_AD__cond_bb2__PB__src_0_wrt_0)
 // CHECK-DATA-STRUCTURES:   case bb1(_AD__cond_bb1__PB__src_0_wrt_0)
 // CHECK-DATA-STRUCTURES: }
-// CHECK-DATA-STRUCTURES: struct _AD__cond_bb3__PB__src_0_wrt_0 {
-// CHECK-DATA-STRUCTURES:   @_hasStorage var predecessor: _AD__cond_bb3__Pred__src_0_wrt_0 { get set }
-// CHECK-DATA-STRUCTURES: }
 
 // CHECK-SIL-LABEL: sil hidden @AD__cond__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {
 // CHECK-SIL: bb0([[INPUT_ARG:%.*]] : $Float):
@@ -137,6 +137,20 @@ func nested_cond_generic<T : Differentiable & FloatingPoint>(_ x: T, _ y: T) ->
   return y
 }
 
+@differentiable
+@_silgen_name("loop_generic")
+func loop_generic<T : Differentiable & FloatingPoint>(_ x: T) -> T {
+  var result = x
+  for _ in 1..<3 {
+    var y = x
+    for _ in 1..<3 {
+      result = y
+      y = result
+    }
+  }
+  return result
+}
+
 // Test control flow + tuple buffer.
 // Verify that adjoint buffers are not allocated for address projections.
 
diff --git a/test/AutoDiff/leakchecking.swift b/test/AutoDiff/leakchecking.swift
@@ -190,6 +190,44 @@ LeakCheckingTests.test("ControlFlow") {
     expectEqual((-2674, 2), Tracked<Float>(-1337).valueWithGradient(in: { x in enum_notactive2(.b(4, 5), x) }))
   }
 
+  // FIXME: Fix control flow AD memory leaks.
+  // See related FIXME comments in adjoint value/buffer propagation in
+  // lib/SILOptimizer/Mandatory/Differentiation.cpp.
+  testWithLeakChecking(expectedLeakCount: 6) {
+    func for_loop(_ x: Tracked<Float>) -> Tracked<Float> {
+      var result = x
+      for _ in 1..<3 {
+        result = result * x
+      }
+      return result
+    }
+    expectEqual((8, 12), Tracked<Float>(2).valueWithGradient(in: for_loop))
+    expectEqual((27, 27), Tracked<Float>(3).valueWithGradient(in: for_loop))
+  }
+
+  // FIXME: Fix control flow AD memory leaks.
+  // See related FIXME comments in adjoint value/buffer propagation in
+  // lib/SILOptimizer/Mandatory/Differentiation.cpp.
+  testWithLeakChecking(expectedLeakCount: 20) {
+    func nested_loop(_ x: Tracked<Float>) -> Tracked<Float> {
+      var outer = x
+      for _ in 1..<3 {
+        outer = outer * x
+
+        var inner = outer
+        var i = 1
+        while i < 3 {
+          inner = inner / x
+          i += 1
+        }
+        outer = inner
+      }
+      return outer
+    }
+    expectEqual((0.5, -0.25), Tracked<Float>(2).valueWithGradient(in: nested_loop))
+    expectEqual((0.25, -0.0625), Tracked<Float>(4).valueWithGradient(in: nested_loop))
+  }
+
   // FIXME: Fix control flow AD memory leaks.
   // See related FIXME comments in adjoint value/buffer propagation in
   // lib/SILOptimizer/Mandatory/Differentiation.cpp.
diff --git a/test/AutoDiff/refcounting.swift b/test/AutoDiff/refcounting.swift
@@ -36,11 +36,11 @@ func testOwnedVector(_ x: Vector) -> Vector {
 }
 _ = pullback(at: Vector.zero, in: testOwnedVector)
 
-// CHECK-LABEL: enum {{.*}}testOwnedVector{{.*}}__Pred__src_0_wrt_0 {
-// CHECK-NEXT: }
 // CHECK-LABEL: struct {{.*}}testOwnedVector{{.*}}__PB__src_0_wrt_0 {
 // CHECK-NEXT:   @_hasStorage var pullback_0: (Vector) -> (Vector, Vector) { get set }
 // CHECK-NEXT: }
+// CHECK-LABEL: enum {{.*}}testOwnedVector{{.*}}__Pred__src_0_wrt_0 {
+// CHECK-NEXT: }
 
 // CHECK-LABEL: sil hidden @{{.*}}UsesMethodOfNoDerivativeMember{{.*}}applied2to{{.*}}__adjoint_src_0_wrt_0_1
 // CHECK: bb0([[SEED:%.*]] : $Vector, [[PB_STRUCT:%.*]] : ${{.*}}UsesMethodOfNoDerivativeMember{{.*}}applied2to{{.*}}__PB__src_0_wrt_0_1):
