Refine debug info emitted for adjoint buffers.

Single input variable might yield multiple adjoint buffers if
control flow is involved. Therefore we cannot simply transfer debug info
from the input variable: it will be invalid as we will end with multiple
locations for a single "source" variable, and, even worse, might end
with conflicting debug info as different buffers might be optimized
differently.

We do:
  - Drop input argument number. This must be unique and we're not
  - Correct variable name

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -559,9 +559,18 @@ class PullbackCloner::Implementation final
     auto debugInfo = findDebugLocationAndVariable(originalValue);
     SILLocation loc = debugInfo ? debugInfo->first.getLocation()
                                 : RegularLocation::getAutoGeneratedLocation();
+    llvm::SmallString<32> adjName;
     auto *newBuf = createFunctionLocalAllocation(
         bufType, loc, /*zeroInitialize*/ true,
-        debugInfo.transform([](AdjointValue::DebugInfo di) { return di.second; }));
+        debugInfo.transform(
+          [&](AdjointValue::DebugInfo di) {
+            SILDebugVariable &dv = di.second;
+            dv.ArgNo = 0;
+            llvm::raw_svector_ostream(adjName) << "adjoint of '" << dv.Name
+                                               << "' in bb" << origBB->getDebugID();
+            dv.Name = adjName;
+            return dv;
+          }));
     return (insertion.first->getSecond() = newBuf);
   }
 

test/AutoDiff/compiler_crashers_fixed/58660-conflicting-debug-info-inlining.swift

@@ -42,9 +42,9 @@ struct MyModel: Differentiable {
   @differentiable(reverse)
   mutating func member4() {
 // CHECK-LABEL: // pullback of MyModel.member4()
-// CHECK-NOT: debug_value %{{.*}} : $MyModel.TangentVector, var, name "self", argno 1, implicit, scope
+// CHECK-NOT: debug_value %{{.*}} : $MyModel.TangentVector, var, name %{{.*}}, argno 1, implicit, scope
 // CHECK: bb1(%{{.*}} : $_AD__$s4main7MyModelV7member4yyF_bb1__PB__src_0_wrt_0):
-// CHECK: debug_value %{{.*}} : $MyModel.TangentVector, var, name "self", argno 1, implicit, loc
+// CHECK: debug_value %{{.*}} : $MyModel.TangentVector, var, name "adjoint of 'self' in bb1", implicit, scope
     // Must be a differentiable type.
     var localVar: Float = 0
 

test/AutoDiff/compiler_crashers_fixed/issue-62608-conflicting-debug-info.swift

@@ -0,0 +1,77 @@
+// RUN: %target-swift-frontend -emit-sil -O -g %s | %FileCheck %s
+
+// Fix for https://github.com/apple/swift/issues/62608
+// We need to emit separate debug info location for different adjoint buffers
+// created for the single input variable
+
+import _Differentiation
+
+public extension Array {
+    @inlinable
+    @differentiable(reverse)
+    mutating func update(at index: Int, byCalling closure: @differentiable(reverse) (inout Element) -> Void) where Element: Differentiable {
+        closure(&self[index])
+    }
+}
+
+public func valueWithPullback<T>(
+    at x: T, of f: @differentiable(reverse) (inout T) -> Void
+) -> (value: Void, pullback: (inout T.TangentVector) -> Void) {
+    @differentiable(reverse)
+    func nonInoutWrappingFunction(_ t: T) -> T {
+        var t = t
+        f(&t)
+        return t
+    }
+    let nonInoutPullback = pullback(at: x, of: nonInoutWrappingFunction)
+    return ((), { $0 = nonInoutPullback($0) })
+}
+
+@inlinable
+public func pullback<T>(
+    at x: T, of f: @differentiable(reverse) (inout T) -> Void
+) -> (inout T.TangentVector) -> Void {
+    return valueWithPullback(at: x, of: f).pullback
+}
+
+// CHECK-LABEL: sil private @$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr :
+// CHECK: alloc_stack $Double, var, name "adjoint of 'element' in bb3"
+// CHECK: debug_value %{{.*}} : $Builtin.FPIEEE64, var, (name "adjoint of 'element' in bb1"
+
+public extension Array where Element: Differentiable {
+    @inlinable
+    @derivative(of: update(at:byCalling:))
+    mutating func vjpUpdate(
+        at index: Int,
+        byCalling closure: @differentiable(reverse) (inout Element) -> Void
+    )
+    ->
+    (value: Void, pullback: (inout Self.TangentVector) -> Void)
+    {
+        let closurePullback = pullback(at: self[index], of: closure)
+        return (value: (), pullback: { closurePullback(&$0.base[index]) })
+    }
+}
+
+func testUpdateByCalling() throws {
+    @differentiable(reverse)
+    func fOfArray(array: [Double]) -> Double {
+        var array = array
+        var result = 0.0
+        for i in withoutDerivative(at: 0 ..< array.count) {
+            array.update(at: i, byCalling: { (element: inout Double) in
+                let initialElement = element
+                for _ in withoutDerivative(at: 0 ..< i) {
+                    element *= initialElement
+                }
+            })
+            result += array[i]
+        }
+        return result
+    }
+    
+    let array = [Double](repeating: 1.0, count: 3)
+    let expectedGradientOfFOfArray = [1.0, 2.0, 3.0]
+    let obtainedGradientOfFOfArray = gradient(at: array, of: fOfArray).base
+}
+