[AutoDiff] Do not propagate same adjoint buffer multiple times

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -35,6 +35,7 @@
 #include "swift/SILOptimizer/PassManager/PrettyStackTrace.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
 
 namespace swift {
 
@@ -2381,9 +2382,11 @@ SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
   assert(pullbackTrampolineBB->getNumArguments() == 1);
   auto loc = origBB->getParent()->getLocation();
   SmallVector<SILValue, 8> trampolineArguments;
+
   // Propagate adjoint values/buffers of active values/buffers to
   // predecessor blocks.
   auto &predBBActiveValues = activeValues[origPredBB];
+  llvm::SmallSet<std::pair<SILValue, SILValue>, 32> propagatedAdjoints;
   for (auto activeValue : predBBActiveValues) {
     LLVM_DEBUG(getADDebugStream()
                << "Propagating adjoint of active value " << activeValue
@@ -2425,12 +2428,14 @@ SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
       // Propagate adjoint buffers using `copy_addr`.
       auto adjBuf = getAdjointBuffer(origBB, activeValue);
       auto predAdjBuf = getAdjointBuffer(origPredBB, activeValue);
-      builder.createCopyAddr(loc, adjBuf, predAdjBuf, IsNotTake,
-                             IsNotInitialization);
+      if (propagatedAdjoints.insert({adjBuf, predAdjBuf}).second)
+        builder.createCopyAddr(loc, adjBuf, predAdjBuf, IsNotTake,
+                               IsNotInitialization);
       break;
     }
     }
   }
+
   // Propagate pullback struct argument.
   TangentBuilder pullbackTrampolineBBBuilder(
       pullbackTrampolineBB, getContext());

test/AutoDiff/SILOptimizer/replicated-adjoint-prop.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -sil-print-after=differentiation -o /dev/null 2>&1 %s | %FileCheck %s -check-prefix=CHECK-SIL
+
+import _Differentiation
+
+struct Test: Differentiable {
+    var val1: Float
+    var val2: Float
+
+    @differentiable(reverse)
+    mutating func doSomething(input: Float) {
+// CHECK-SIL-LABEL: sil private [ossa] @$s4null4TestV11doSomething5inputySf_tFTJpSSpSr : $@convention(thin)
+// Ensure that only two adjoint buffers will be propagated    
+// CHECK-SIL: copy_addr %0 to %22 : $*Test.TangentVector
+// CHECK-SIL-NEXT: debug_value
+// CHECK-SIL-NEXT: copy_addr %0 to %18 : $*Test.TangentVector
+// CHECK-SIL-NEXT: switch_enum %1
+        self.val1 *= input
+        self.val2 *= input
+
+        if self.val1 > input {
+            self.val1 = input
+        }
+        if self.val2 > input {
+            self.val2 = input
+        }
+    }
+}
+
+@differentiable(reverse)
+func wrapper(input: Float, multiplier: Float) -> Float {
+    var test = Test(val1: input, val2: input)
+    test.doSomething(input: multiplier)
+    return test.val1 * test.val2
+}
+
+let grad = gradient(at: 2.0, 3.0, of: wrapper)
+print("Grad: \(grad)")