Incorrect gradients in some optional-related cases

[asl]

Description

The code below should produce identical gradients, however, it does not

Reproduction

import _Differentiation

@differentiable(reverse)
func square(y: Double) -> Double? {
  return y * y
}

@differentiable(reverse)
func square2(y: Double?) -> Double {
  return y! * y!
}

print(valueWithGradient(at: 42.0, of: {y in square(y:y)!}))
print(valueWithGradient(at: 42.0, of: {y in square2(y:y)}))

Expected behavior

(value: 1764.0, gradient: 84.0)
(value: 1764.0, gradient: 84.0)

Environment

Swift version 6.0-dev (LLVM 76ba61e88550d3b, Swift 9470ad7)
Target: arm64-apple-macosx13.0

Additional information

Currently we are printing:
(value: 1764.0, gradient: 84.0)
(value: 1764.0, gradient: 42.0)

[asl]

Tagging @rxwei

[asl]

The following code produces correct result:

@differentiable(reverse)
func square2(y: Double?) -> Double {
  let some = y!
  return some * some
}

It looks like some optional adjoint buffer accumulation is wrong.

[asl]

@rxwei I think non-switch_enum case is broken as well, though I did not yet come with the testcase that could show the incorrect results. All cases I can cook seem to be equivalent one way or another :)

The code in question is https://github.com/swiftlang/swift/pull/74985/files#diff-01cf87f81a8c47d84be8508fac2cb0f1a4ba15919ad86d3c8904bccee60151b5L2967

Here we are always overwriting adjoint of the incoming value in the predecessor block. However, is it possible that incoming value has multiple uses? As a result, in the reverse pass we'll have multiple adjoint values that should be accumulated?

The address-only case below does the accumulation in the adjoint buffer – I guess it is just much easier to obtain testcase here via some generics.

So, essentially I'm thinking about the following patch:

diff --git a/lib/SILOptimizer/Differentiation/PullbackCloner.cpp b/lib/SILOptimizer/Differentiation/PullbackCloner.cpp
index d276fdde8ed..f4e8b0da0c7 100644
--- a/lib/SILOptimizer/Differentiation/PullbackCloner.cpp
+++ b/lib/SILOptimizer/Differentiation/PullbackCloner.cpp
@@ -2966,8 +2966,8 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
           } else {
             blockTemporaries[getPullbackBlock(predBB)].insert(
               concreteBBArgAdjCopy);
-            setAdjointValue(predBB, incomingValue,
-                            makeConcreteAdjointValue(concreteBBArgAdjCopy));
+            addAdjointValue(predBB, incomingValue,
+                            makeConcreteAdjointValue(concreteBBArgAdjCopy), pbLoc);
           }
         }
         break;

It seems to be safe to me:

• If there was no adjoint for incomingValue registered yet, then it's fine, it is equivalent to setAdjointValue
• If there was some adjoint in predBB known for incomingValue it would perform necessary accumulation (this is only possible it seems if there are multiple active uses of incomingValue).

What do you think?