[AutoDiff] Fixes small bugs in the PR fixing #66522

During internal testing we discovered 2 more bugs -
1. The element adjoint of a struct_extract can itself be an AddElement.
2. Indirect concrete adjoint materialization was missing a copy operation.

This commit fixes these bugs and adds relevant test cases.

Author: jkshtj

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -424,7 +424,8 @@ class PullbackCloner::Implementation final
     /// the destination address.
     case AdjointValueKind::Concrete: {
       auto concreteVal = val.getConcreteValue();
-      builder.emitStoreValueOperation(loc, concreteVal, destAddress,
+      auto copyOfConcreteVal = builder.emitCopyValueOperation(loc, concreteVal);
+      builder.emitStoreValueOperation(loc, copyOfConcreteVal, destAddress,
                                       StoreOwnershipQualifier::Init);
       break;
     }
@@ -1241,17 +1242,14 @@ class PullbackCloner::Implementation final
         break;
       }
       case AdjointValueKind::Aggregate:
-      case AdjointValueKind::Concrete: {
+      case AdjointValueKind::Concrete:
+      case AdjointValueKind::AddElement: {
         auto baseAdj = makeZeroAdjointValue(tangentVectorSILTy);
         addAdjointValue(bb, sei->getOperand(),
                         makeAddElementAdjointValue(baseAdj, eltAdj, tanField),
                         loc);
         break;
       }
-      case AdjointValueKind::AddElement: {
-        llvm_unreachable(
-            "Adjoint of extracted element in `StructExtractInst` cannot be of kind `AddElement`");
-      }
       }
       break;
     }
@@ -1439,7 +1437,8 @@ class PullbackCloner::Implementation final
       break;
     }
     case AdjointValueKind::Aggregate:
-    case AdjointValueKind::Concrete: {
+    case AdjointValueKind::Concrete:
+    case AdjointValueKind::AddElement: {
       auto tupleTy = tei->getTupleType();
       auto tupleTanTupleTy = tupleTanTy.getAs<TupleType>();
       if (!tupleTanTupleTy) {
@@ -1466,10 +1465,6 @@ class PullbackCloner::Implementation final
       }
       break;
     }
-    case AdjointValueKind::AddElement: {
-        llvm_unreachable(
-            "Adjoint of extracted element in `TupleExtractInst` cannot be of kind `AddElement`");
-    }
     }
   }
 

test/AutoDiff/SILOptimizer/pullback_generation.sil

@@ -124,9 +124,13 @@ bb0(%0 : $(Float, Float)):
 // CHECK: } // end sil function 'function_with_tuple_extract_1TJpSpSr'
 
 //===----------------------------------------------------------------------===//
-// Pullback generation - `tuple_extract`
-// - Input to pullback has non-owned ownership semantics which requires copying
-// this value to stack before lifetime-ending uses. 
+// Pullback generation - Inner values of concrete adjoints must be copied 
+// during direct materialization. 
+// - If the input to pullback BB has non-owned ownership semantics we cannot 
+// perform a lifetime-ending operation on it.
+// - If the input to the pullback BB is an owned, non-trivial value we must 
+// copy it or there will be a double consume when all owned parameters are 
+// destroyed at the end of the basic block.
 //===----------------------------------------------------------------------===//
 sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @function_with_tuple_extract_2: $@convention(thin) (@guaranteed (X, X)) -> @owned X {
 }
@@ -164,3 +168,19 @@ bb0(%0 : @guaranteed $(X, X)):
 // CHECK:   destroy_value %19 : $(X, X)                     
 // CHECK:   return %21 : $(X, X)                            
 // CHECK: } // end sil function 'function_with_tuple_extract_2TJpSpSr'
+
+//===----------------------------------------------------------------------===//
+// Pullback generation - `tuple_extract`
+// - Adjoint of extracted element can be `AddElement`
+// - Just need to make sure that we are able to generate a pullback
+//===----------------------------------------------------------------------===//
+sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @function_with_tuple_extract_3: $@convention(thin) (((Float, Float), Float)) -> Float {
+}
+
+sil hidden [ossa] @function_with_tuple_extract_3: $@convention(thin) (((Float, Float), Float)) -> Float {
+bb0(%0 : $((Float, Float), Float)):
+  %1 = tuple_extract %0 : $((Float, Float), Float), 0
+  %2 = tuple_extract %1 : $(Float, Float), 0
+  return %2 : $Float
+}
+// CHECK-LABEL: sil private [ossa] @function_with_tuple_extract_3TJpSpSr : $@convention(thin) (Float) -> ((Float, Float), Float) {

test/AutoDiff/SILOptimizer/pullback_generation.swift

@@ -128,3 +128,73 @@ func f3(v: NonPiecewiseMaterializableWithAggDifferentiableField) -> PiecewiseMat
 // CHECK:   destroy_value %13 : $NonPiecewiseMaterializableWithAggDifferentiableField 
 // CHECK:   return %16 : $NonPiecewiseMaterializableWithAggDifferentiableField 
 // CHECK: } // end sil function '$s19pullback_generation2f31vAA23PiecewiseMaterializableVAA03NondE26WithAggDifferentiableFieldV_tFTJpSpSr'
+
+//===----------------------------------------------------------------------===//
+// Pullback generation - `struct_extract`
+// - Adjoint of extracted element can be `AddElement`
+// - Just need to make sure that we are able to generate a pullback for B.x's 
+// getter
+//===----------------------------------------------------------------------===//
+struct A: Differentiable {
+    public var x: Float
+}
+
+struct B: Differentiable {
+    var y: A
+    
+    public init(a: A) {
+        self.y = a
+    }
+    
+    @differentiable(reverse)
+    public var x: Float {
+        get { return self.y.x }
+    }
+}
+
+// CHECK-LABEL: sil private [ossa] @$s19pullback_generation1BV1xSfvgTJpSpSr : $@convention(thin) (Float) -> B.TangentVector {
+
+//===----------------------------------------------------------------------===//
+// Pullback generation - Inner values of concrete adjoints must be copied 
+// during indirect materialization
+//===----------------------------------------------------------------------===//
+
+struct NonTrivial {
+    var x: Float
+    var y: String
+}
+
+extension NonTrivial: Differentiable, Equatable, AdditiveArithmetic {
+    public typealias TangentVector = Self
+    mutating func move(by offset: TangentVector) {fatalError()}
+    public static var zero: Self {fatalError()}
+    public static func + (lhs: Self, rhs: Self) -> Self {fatalError()}
+    public static func - (lhs: Self, rhs: Self) -> Self {fatalError()}
+}
+
+@differentiable(reverse)
+func f4(a: NonTrivial) -> Float {
+    var sum: Float = 0
+    for _ in 0..<1 {
+        sum += a.x
+    }
+    return sum
+}
+
+// CHECK-LABEL: sil private [ossa] @$s19pullback_generation2f41aSfAA10NonTrivialV_tFTJpSpSr : $@convention(thin) (Float, @guaranteed Builtin.NativeObject) -> @owned NonTrivial {
+// CHECK: bb5(%67 : @owned $NonTrivial, %68 : $Float, %69 : @owned $(predecessor: _AD__$s19pullback_generation2f41aSfAA10NonTrivialV_tF_bb2__Pred__src_0_wrt_0, @callee_guaranteed (@inout Float) -> Float)):
+// CHECK: %88 = alloc_stack $NonTrivial
+
+// Non-trivial value must be copied or there will be a
+// double consume when all owned parameters are destroyed 
+// at the end of the basic block.
+// CHECK: %89 = copy_value %67 : $NonTrivial              
+
+// CHECK: store %89 to [init] %88 : $*NonTrivial          
+// CHECK: %91 = struct_element_addr %88 : $*NonTrivial, #NonTrivial.x 
+// CHECK: %92 = alloc_stack $Float                        
+// CHECK: store %86 to [trivial] %92 : $*Float            
+// CHECK: %94 = witness_method $Float, #AdditiveArithmetic."+=" : <Self where Self : AdditiveArithmetic> (Self.Type) -> (inout Self, Self) -> () : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@inout τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> () 
+// CHECK: %95 = metatype $@thick Float.Type               
+// CHECK: %96 = apply %94<Float>(%91, %92, %95) : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@inout τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> ()
+// CHECK: destroy_value %67 : $NonTrivial