[AutoDiff] Fix aggregate adjoint value accumulation bug. (#28477)

`PullbackEmitter::accumulateAdjointsDirect` did not handle the case where
`lhs` is an `Aggregate` value and `rhs` is a `Concrete` value.

Resolves TF-943.

Author: bgogul

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -7740,8 +7740,9 @@ PullbackEmitter::accumulateAdjointsDirect(AdjointValue lhs, AdjointValue rhs,
   // (x, y)
   case AdjointValueKind::Aggregate:
     switch (rhs.getKind()) {
-    // (x, y) + z => (x + z.0, y + z.1)
+    // (x, y) + z => (z.0 + x, z.1 + y)
     case AdjointValueKind::Concrete:
+      return accumulateAdjointsDirect(rhs, lhs, loc);
     // x + 0 => x
     case AdjointValueKind::Zero:
       return lhs;

test/AutoDiff/simple_math.swift

@@ -1,3 +1,4 @@
+// RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null 2>&1 | %FileCheck %s
 // RUN: %target-run-simple-swift
 // NOTE(TF-813): verify that enabling forward-mode does not affect reverse-mode.
 // RUN: %target_run_simple_swift_forward_mode_differentiation
@@ -203,7 +204,7 @@ SimpleMathTests.test("StructMemberwiseInitializer") {
     let foo = Foo(stored: input)
     return foo.computed * foo.stored
   }
-  expectEqual(16, 𝛁product)
+  expectEqual(48, 𝛁product)
 
   struct Custom : AdditiveArithmetic, Differentiable {
     var x: Float
@@ -350,4 +351,37 @@ SimpleMathTests.test("ForceUnwrapping") {
   expectEqual((1, 2), forceUnwrap(Float(2)))
 }
 
+// CHECK-LABEL: sil hidden [ossa] @AD__${{.*}}jumpTimesTwo{{.*}}pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned _AD__$s4nullyycfU18_12jumpTimesTwoL_5modelSfAAyycfU18_14SmallTestModelL_V_tF_bb0__PB__src_0_wrt_0) -> SmallTestModel.TangentVector {
+// CHECK: bb0([[DX:%.*]] : $Float,  [[PB_STRUCT:%.*]] : {{.*}}):
+// CHECK:   ([[PB0:%.*]], [[PB1:%.*]]) = destructure_struct [[PB_STRUCT]]
+// CHECK:   [[ADJ_TUPLE:%.*]] = apply [[PB1]]([[DX]]) : $@callee_guaranteed (Float) -> (Float, Float)
+// CHECK:   ([[TMP0:%.*]], [[ADJ_CONCRETE:%.*]]) = destructure_tuple [[ADJ_TUPLE]] : $(Float, Float)
+// CHECK:   [[TMP1:%.*]] = apply [[PB0]]([[TMP0]]) : $@callee_guaranteed (Float) -> SmallTestModel.TangentVector
+// CHECK:   [[ADJ_STRUCT_FIELD:%.*]] = destructure_struct [[TMP1]] : $SmallTestModel.TangentVector
+// CHECK:   [[TMP_RES:%.*]] = alloc_stack $Float
+// CHECK:   [[TMP_ADJ_STRUCT_FIELD:%.*]] = alloc_stack $Float
+// CHECK:   [[TMP_ADJ_CONCRETE:%.*]] = alloc_stack $Float
+// CHECK:   store [[ADJ_STRUCT_FIELD]] to [trivial] [[TMP_ADJ_STRUCT_FIELD]] : $*Float
+// CHECK:   store [[ADJ_CONCRETE]] to [trivial] [[TMP_ADJ_CONCRETE]] : $*Float
+// CHECK:   [[PLUS_EQUAL:%.*]] = witness_method $Float, #AdditiveArithmetic."+"
+// CHECK:   %{{.*}} = apply [[PLUS_EQUAL]]<Float>([[TMP_RES]], [[TMP_ADJ_CONCRETE]], [[TMP_ADJ_STRUCT_FIELD]], {{.*}})
+// CHECK:   [[RES:%.*]] = load [trivial] [[TMP_RES]] : $*Float
+// CHECK:   [[RES_STRUCT:%.*]] = struct $SmallTestModel.TangentVector ([[RES]] : $Float)
+// CHECK:   return [[RES_STRUCT]] : $SmallTestModel.TangentVector
+// CHECK: }
+
+SimpleMathTests.test("Struct") {
+  // TF-943: Test adjoint value accumulation for aggregate lhs and concrete rhs.
+  struct SmallTestModel : Differentiable {
+    public var jump: Float = 3.0
+    @differentiable public func callAsFunction() -> Float { return jump }
+  }
+
+  func jumpTimesTwo(model: SmallTestModel) -> Float{
+    return model() + model.jump
+  }
+  let grads = gradient(at: SmallTestModel(), in: jumpTimesTwo)
+  expectEqual(2.0, grads.jump)
+}
+
 runAllTests()