[AutoDiff] Refine debug info emitted for adjoint buffers (#62779)

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,23 @@ 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) {
+            llvm::raw_svector_ostream adjNameStream(adjName);
+            SILDebugVariable &dv = di.second;
+            dv.ArgNo = 0;
+            adjNameStream << "derivative of '" << dv.Name << "'";
+            if (SILDebugLocation origBBLoc = origBB->front().getDebugLocation()) {
+              adjNameStream << " in scope at ";
+              origBBLoc.getLocation().print(adjNameStream, getASTContext().SourceMgr);
+            }
+            adjNameStream << " (scope #" << origBB->getDebugID() << ")";
+            dv.Name = adjName;
+            return dv;
+          }));
     return (insertion.first->getSecond() = newBuf);
   }
 

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

@@ -1,6 +1,8 @@
 // RUN: %target-build-swift %s
 // RUN: %target-swift-frontend -emit-sil -O -g %s | %FileCheck %s
 
+// REQUIRES: swift_in_compiler
+
 // Issue #58660: Specifically-shaped differentiable functions yield "conflicting debug info for argument" assertion failure
 // Ensure that proper location is preserved after sil-mem2reg location-less stores (created during inlining)
 
@@ -42,9 +44,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 "derivative of 'self' in scope at {{.*}} (scope #1)", 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,79 @@
+// RUN: %target-swift-frontend -emit-sil -O -g %s | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+
+// 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 "derivative of 'element' in scope at {{.*}} (scope #3)"
+// CHECK: debug_value %{{.*}} : $Builtin.FPIEEE64, var, (name "derivative of 'element' in scope at {{.*}} (scope #1)"
+
+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
+}
+