Attempt at a patch for SR-14053 and SR-14218, based on Dan Zheng's initial fix.

Author: BradLarson

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -2052,6 +2052,8 @@ bool PullbackCloner::Implementation::run() {
   SmallVector<SILValue, 8> retElts;
   // This vector will contain all indirect parameter adjoint buffers.
   SmallVector<SILValue, 4> indParamAdjoints;
+  // This vector will identify the locations where initialization is needed.
+  SmallVector<bool, 8> outputsToInitialize;
 
   auto conv = getOriginal().getConventions();
   auto origParams = getOriginal().getArgumentsWithoutIndirectResults();
@@ -2071,27 +2073,62 @@ bool PullbackCloner::Implementation::run() {
     case SILValueCategory::Address: {
       auto adjBuf = getAdjointBuffer(origEntry, origParam);
       indParamAdjoints.push_back(adjBuf);
+      outputsToInitialize.push_back(
+        !conv.getParameters()[parameterIndex].isIndirectMutating());
       break;
     }
     }
   };
+  SmallVector<SILArgument *, 4> pullbackIndirectResults(
+        getPullback().getIndirectResults().begin(),
+        getPullback().getIndirectResults().end());
+
   // Collect differentiation parameter adjoints.
+  // Do a first pass to collect non-inout values.
+  unsigned pullbackInoutArgumentIndex = 0;
   for (auto i : getConfig().parameterIndices->getIndices()) {
-    // Skip `inout` parameters.
-    if (conv.getParameters()[i].isIndirectMutating())
-      continue;
-    addRetElt(i);
+    auto isParameterInout = conv.getParameters()[i].isIndirectMutating();
+    if (!isParameterInout) {
+      addRetElt(i);
+    }
+  }
+
+  // Do a second pass for all inout values.
+  for (auto i : getConfig().parameterIndices->getIndices()) {
+    // Skip `inout` parameters for functions with single basic blocks:
+    // additional adjoint accumulation is not necessary.
+
+    // For functions with multiple basic blocks, accumulation is needed
+    // for `inout` parameters because pullback basic blocks have different
+    // adjoint buffers.
+    auto isOriginalFunctionSingleBasicBlock = getOriginal().size() == 1;
+    auto isParameterInout = conv.getParameters()[i].isIndirectMutating();
+    if (isParameterInout) {
+      if (isOriginalFunctionSingleBasicBlock) {
+        continue;
+      } else {
+        pullbackIndirectResults.push_back(
+          getPullback().getArgumentsWithoutIndirectResults()[pullbackInoutArgumentIndex++]);
+      }
+      addRetElt(i);
+    }
   }
 
   // Copy them to adjoint indirect results.
-  assert(indParamAdjoints.size() == getPullback().getIndirectResults().size() &&
+  assert(indParamAdjoints.size() == pullbackIndirectResults.size() &&
          "Indirect parameter adjoint count mismatch");
-  for (auto pair : zip(indParamAdjoints, getPullback().getIndirectResults())) {
+  unsigned currentIndex = 0;
+  for (auto pair : zip(indParamAdjoints, pullbackIndirectResults)) {
     auto source = std::get<0>(pair);
     auto *dest = std::get<1>(pair);
-    builder.createCopyAddr(pbLoc, source, dest, IsTake, IsInitialization);
+    if (outputsToInitialize[currentIndex]) {
+      builder.createCopyAddr(pbLoc, source, dest, IsTake, IsInitialization);
+    } else {
+      builder.createCopyAddr(pbLoc, source, dest, IsTake, IsNotInitialization);
+    }
     // Prevent source buffer from being deallocated, since the underlying
     // value is moved.
+    currentIndex++;
     destroyedLocalAllocations.insert(source);
   }
 

test/AutoDiff/validation-test/inout_control_flow.swift

@@ -0,0 +1,96 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+import _Differentiation
+
+var InoutControlFlowTests = TestSuite("InoutControlFlow")
+
+// SR-14218
+struct Model: Differentiable {
+  var first: Float = 3
+  var second: Float = 1
+
+  mutating func outer(){
+    inner()
+  }
+
+  mutating func inner() {
+    self.second = self.first
+
+    let x = 5
+    if x < 50 {}
+  }
+}
+
+@differentiable(reverse)
+func loss(model: Model) -> Float{
+  var model = model
+  model.outer()
+  return model.second
+}
+
+InoutControlFlowTests.test("MutatingBeforeControlFlow") {
+  var model = Model()
+  let grad = gradient(at: model, of: loss)
+  expectEqual(1, grad.first)
+  expectEqual(0, grad.second)
+}
+
+// SR-14053
+protocol NumericDifferentiable : Numeric, Differentiable {
+  @differentiable(reverse) static func *(lhs: Self, rhs: Self) -> Self
+}
+
+extension Float: NumericDifferentiable {}
+
+struct Model2<T: NumericDifferentiable>: Differentiable {
+  var first: T
+  var second: T
+}
+
+@differentiable(reverse)
+func adjust<T: NumericDifferentiable>(model: inout Model2<T>, multiplier: T) {
+  model.first = model.second * multiplier
+
+  let x = 5
+  if x < 50 {}
+}
+
+@differentiable(reverse)
+func loss2(model: Model2<Float>, multiplier: Float) -> Float {
+  var model = model
+  adjust(model: &model, multiplier: multiplier)
+  return model.first
+}
+
+InoutControlFlowTests.test("InoutParameterWithControlFlow") {
+  var model = Model2<Float>(first: 1, second: 3)
+  let grad = gradient(at: model, 5.0, of: loss2)
+  expectEqual(0, grad.0.first)
+  expectEqual(5, grad.0.second)
+}
+
+@differentiable(reverse)
+func adjust2<T: NumericDifferentiable>(multiplier: T, model: inout Model2<T>) {
+  model.first = model.second * multiplier
+
+  let x = 5
+  if x < 50 {}
+}
+
+@differentiable(reverse)
+func loss3(model: Model2<Float>, multiplier: Float) -> Float {
+  var model = model
+  adjust2(multiplier: multiplier, model: &model)
+  return model.first
+}
+
+InoutControlFlowTests.test("LaterInoutParameterWithControlFlow") {
+  var model = Model2<Float>(first: 1, second: 3)
+  let grad = gradient(at: model, 5.0, of: loss3)
+  expectEqual(0, grad.0.first)
+  expectEqual(5, grad.0.second)
+}
+
+runAllTests()