[AutoDiff] Do not propagate same adjoint buffer multiple times (#64963)

Adjoint buffers of projections (e.g. obtained via begin_access) are same as adjoint buffer of underlying struct value. As a result, when propagating adjoint values to pullback successor blocks we tend to produce lots of identical copies (essentially for every struct access and in every basic block) of adjoint buffers.

These copy_addrs instructions are then lowered down to plain loads and stores and while the redundant copies are usually optimized away by subsequent optimization passes, presence of such copies leads to elevated memory consumption and compilation time as one needs to track liveness of these values being copied.

Track the values being propagated and simply do not generate extra copies if the same value was already propagated.

One step towards #61773

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: TestV11doSomething5inputySf_tFTJpSSpSr :
+// 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)")