[AutoDiff] Correct propagate adjoints for array literal values (#81676)

Fixes #81607

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -3666,14 +3666,31 @@ void PullbackCloner::Implementation::
   if (originalValue != dti->getResult(0))
     return;
   // Accumulate the array's adjoint value into the adjoint buffers of its
-  // element addresses: `pointer_to_address` and `index_addr` instructions.
+  // element addresses: `pointer_to_address` and (optionally) `index_addr`
+  // instructions.
+  // The input code looks like as follows:
+  //  %17 = integer_literal $Builtin.Word, 1
+  //  function_ref _allocateUninitializedArray<A>(_:)
+  //  %18 = function_ref @$ss27_allocateUninitializedArrayySayxG_BptBwlF : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
+  //  %19 = apply %18<Float>(%17) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
+  //  (%20, %21) = destructure_tuple %19
+  //  %22 = mark_dependence %21 on %20
+  //  %23 = pointer_to_address %22 to [strict] $*Float
+  //  store %0 to [trivial] %23
+  //  function_ref _finalizeUninitializedArray<A>(_:)
+  //  %25 = function_ref @$ss27_finalizeUninitializedArrayySayxGABnlF : $@convention(thin) <τ_0_0> (@owned Array<τ_0_0>) -> @owned Array<τ_0_0>
+  //  %26 = apply %25<Float>(%20) : $@convention(thin) <τ_0_0> (@owned Array<τ_0_0>) -> @owned Array<τ_0_0> // user: %27
+  // Note that %20 and %21 are in some sense "aliases" for each other. Here our `originalValue` is %20 in the code above.
+  // We need to trace from %21 down to %23 and propagate (decomposed) adjoint of originalValue to adjoint of %23.
+  // Then the generic adjoint propagation code would do its job to propagate %23' to %0'.
+  // If we're initializing multiple values we're having additional `index_addr` instructions, but
+  // the handling is similar.
   LLVM_DEBUG(getADDebugStream()
              << "Accumulating adjoint value for array literal into element "
                 "address adjoint buffers"
              << originalValue);
   auto arrayAdjoint = materializeAdjointDirect(arrayAdjointValue, loc);
   builder.setCurrentDebugScope(remapScope(dti->getDebugScope()));
-  builder.setInsertionPoint(arrayAdjoint->getParentBlock());
   for (auto use : dti->getResult(1)->getUses()) {
     auto *mdi = dyn_cast<MarkDependenceInst>(use->getUser());
     assert(mdi && "Expected mark_dependence user");

test/AutoDiff/compiler_crashers_fixed/issue-81607-array-literal-adjoint.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// https://github.com/swiftlang/swift/issues/81607
+// Ensure we're propagating array adjoint in the correct BB
+
+import _Differentiation
+
+@differentiable(reverse)
+func sum(_ a: Float, _ b: [Float]) -> [Float] {
+  if b.count != 0 {
+    return b
+  }
+  return [a]
+}