Incorporated Dan's cleanup suggestions.

Author: BradLarson

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -2093,25 +2093,27 @@ bool PullbackCloner::Implementation::run() {
     }
   }
 
-  // Do a second pass for all inout values.
+  // Do a second pass for all inout parameters.
   for (auto i : getConfig().parameterIndices->getIndices()) {
-    // Skip `inout` parameters for functions with single basic blocks:
-    // additional adjoint accumulation is not necessary.
+    // Skip non-inout parameters.
+    auto isParameterInout = conv.getParameters()[i].isIndirectMutating();
+    if (!isParameterInout)
+      continue;
+
+    // Skip `inout` parameters for functions with a single basic block:
+    // adjoint accumulation for those parameters is already done by
+    // per-instruction visitors.
+    if (getOriginal().size() == 1)
+      continue;
 
     // 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);
-    }
+    auto pullbackInoutArgument =
+        getPullback()
+            .getArgumentsWithoutIndirectResults()[pullbackInoutArgumentIndex++];
+    pullbackIndirectResults.push_back(pullbackInoutArgument);
+    addRetElt(i);
   }
 
   // Copy them to adjoint indirect results.
@@ -2126,9 +2128,9 @@ bool PullbackCloner::Implementation::run() {
     } else {
       builder.createCopyAddr(pbLoc, source, dest, IsTake, IsNotInitialization);
     }
+    currentIndex++;
     // Prevent source buffer from being deallocated, since the underlying
     // value is moved.
-    currentIndex++;
     destroyedLocalAllocations.insert(source);
   }
 

test/AutoDiff/validation-test/inout_control_flow.swift

@@ -11,13 +11,14 @@ struct Model: Differentiable {
   var first: Float = 3
   var second: Float = 1
 
-  mutating func outer(){
+  mutating func outer() {
     inner()
   }
 
   mutating func inner() {
     self.second = self.first
 
+    // Dummy no-op if block, required to introduce control flow.
     let x = 5
     if x < 50 {}
   }
@@ -53,6 +54,7 @@ struct Model2<T: NumericDifferentiable>: Differentiable {
 func adjust<T: NumericDifferentiable>(model: inout Model2<T>, multiplier: T) {
   model.first = model.second * multiplier
 
+  // Dummy no-op if block, required to introduce control flow.
   let x = 5
   if x < 50 {}
 }
@@ -75,6 +77,7 @@ InoutControlFlowTests.test("InoutParameterWithControlFlow") {
 func adjust2<T: NumericDifferentiable>(multiplier: T, model: inout Model2<T>) {
   model.first = model.second * multiplier
 
+  // Dummy no-op if block, required to introduce control flow.
   let x = 5
   if x < 50 {}
 }