[AutoDiff] Support ref_element_addr differentiation. (#29749)

Support differentiation of `ref_element_addr`: class stored property references.

Activity analysis:
- Propagate activity for `ref_element_addr` like `struct_element_addr`.
- Do not propagate activity for `ref_element_addr` to `@noDerivative` members.

Pullback generation rules:

```
  Original: y = ref_element_addr x, <n>
   Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
                                      ^~~~~~~
                    field in tangent space corresponding to #field
```

- Add `ref_element_addr` case to `PullbackEmitter::getAdjointProjection`.
  - The adjoint projection of a `ref_element_addr` is a local allocation
    initialized with the corresponding field value from the class's base
    adjoint value.

Exposes TF-1149: cannot differentiate active value with loadable type but
address-only `TangentVector` type. Diagnose for now.

Resolves SR-12152.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -491,6 +491,12 @@ WARNING(autodiff_nonvaried_result_fixit,none,
         "result does not depend on differentiation arguments and will always "
         "have a zero derivative; do you want to use 'withoutDerivative(at:)'?",
         ())
+// TODO(TF-1149): Remove this diagnostic.
+NOTE(autodiff_loadable_value_addressonly_tangent_unsupported,none,
+     "cannot yet differentiate value whose type %0 has a compile-time known "
+     "size, but whose 'TangentVector' contains stored properties of unknown "
+     "size; consider modifying %1 to use fewer generic parameters in stored "
+     "properties", (Type, Type))
 NOTE(autodiff_enums_unsupported,none,
      "differentiating enum values is not yet supported", ())
 NOTE(autodiff_global_let_closure_not_differentiable,none,
@@ -543,9 +549,6 @@ NOTE(autodiff_jvp_control_flow_not_supported,none,
      "forward-mode differentiation does not yet support control flow", ())
 NOTE(autodiff_control_flow_not_supported,none,
      "cannot differentiate unsupported control flow", ())
-// TODO(TF-645): Remove when differentiation supports `ref_element_addr`.
-NOTE(autodiff_class_property_not_supported,none,
-     "differentiating class properties is not yet supported", ())
 // TODO(TF-1080): Remove when differentiation supports `begin_apply`.
 NOTE(autodiff_coroutines_not_supported,none,
      "differentiation of coroutine calls is not yet supported", ())

include/swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h

@@ -157,8 +157,8 @@ class DifferentiableActivityInfo {
   /// Propagates variedness from the given operand to its user's results.
   void propagateVaried(Operand *operand, unsigned independentVariableIndex);
   /// Marks the given value as varied and recursively propagates variedness
-  /// inwards (to operands) through projections. Skips `@noDerivative` struct
-  /// field projections.
+  /// inwards (to operands) through projections. Skips `@noDerivative` field
+  /// projections.
   void
   propagateVariedInwardsThroughProjections(SILValue value,
                                            unsigned independentVariableIndex);
@@ -172,9 +172,9 @@ class DifferentiableActivityInfo {
                                        unsigned dependentVariableIndex);
   /// Propagates usefulnesss to the operands of the given instruction.
   void propagateUseful(SILInstruction *inst, unsigned dependentVariableIndex);
-  /// Marks the given address as useful and recursively propagates usefulness
-  /// inwards (to operands) through projections. Skips `@noDerivative` struct
-  /// field projections.
+  /// Marks the given address or class-typed value as useful and recursively
+  /// propagates usefulness inwards (to operands) through projections. Skips
+  /// `@noDerivative` field projections.
   void propagateUsefulThroughAddress(SILValue value,
                                      unsigned dependentVariableIndex);
   /// If the given value is an `array.uninitialized_intrinsic` application,

include/swift/SILOptimizer/Utils/Differentiation/PullbackEmitter.h

@@ -266,10 +266,17 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
 
   SILBasicBlock::iterator getNextFunctionLocalAllocationInsertionPoint();
 
+  /// Creates and returns a local allocation with the given type.
+  ///
+  /// Local allocations are created uninitialized in the pullback entry and
+  /// deallocated in the pullback exit. All local allocations not in
+  /// `destroyedLocalAllocations` are also destroyed in the pullback exit.
+  AllocStackInst *createFunctionLocalAllocation(SILType type, SILLocation loc);
+
   SILValue &getAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer);
 
-  // Accumulates `rhsBufferAccess` into the adjoint buffer corresponding to
-  // `originalBuffer`.
+  /// Accumulates `rhsBufferAccess` into the adjoint buffer corresponding to
+  /// `originalBuffer`.
   void addToAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer,
                           SILValue rhsBufferAccess, SILLocation loc);
 
@@ -354,6 +361,13 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   ///                     field in tangent space corresponding to #field
   void visitStructExtractInst(StructExtractInst *sei);
 
+  /// Handle `ref_element_addr` instruction.
+  ///   Original: y = ref_element_addr x, <n>
+  ///    Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
+  ///                                       ^~~~~~~
+  ///                     field in tangent space corresponding to #field
+  void visitRefElementAddrInst(RefElementAddrInst *reai);
+
   /// Handle `tuple` instruction.
   ///   Original: y = tuple (x0, x1, x2, ...)
   ///    Adjoint: (adj[x0], adj[x1], adj[x2], ...) += destructure_tuple adj[y]
@@ -421,8 +435,6 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
       UnconditionalCheckedCastAddrInst *uccai);
 
 #define NOT_DIFFERENTIABLE(INST, DIAG) void visit##INST##Inst(INST##Inst *inst);
-
-  NOT_DIFFERENTIABLE(RefElementAddr, autodiff_class_property_not_supported)
 #undef NOT_DIFFERENTIABLE
 
 #define NO_ADJOINT(INST)                                                       \

lib/SIL/SILFunctionType.cpp

@@ -275,9 +275,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     switch (origResConv) {
     case ResultConvention::Owned:
     case ResultConvention::Autoreleased:
-      conv = tl.isTrivial()
-          ? ParameterConvention::Direct_Unowned
-          : ParameterConvention::Direct_Guaranteed;
+      if (tl.isAddressOnly()) {
+        conv = ParameterConvention::Indirect_In_Guaranteed;
+      } else {
+        conv = tl.isTrivial() ? ParameterConvention::Direct_Unowned
+                              : ParameterConvention::Direct_Guaranteed;
+      }
       break;
     case ResultConvention::Unowned:
     case ResultConvention::UnownedInnerPointer:
@@ -301,9 +304,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     case ParameterConvention::Direct_Owned:
     case ParameterConvention::Direct_Guaranteed:
     case ParameterConvention::Direct_Unowned:
-      conv = tl.isTrivial()
-          ? ResultConvention::Unowned
-          : ResultConvention::Owned;
+      if (tl.isAddressOnly()) {
+        conv = ResultConvention::Indirect;
+      } else {
+        conv = tl.isTrivial() ? ResultConvention::Unowned
+                              : ResultConvention::Owned;
+      }
       break;
     case ParameterConvention::Indirect_In:
     case ParameterConvention::Indirect_Inout:

lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp

@@ -166,18 +166,20 @@ void DifferentiableActivityInfo::propagateVaried(
         setVariedAndPropagateToUsers(teai, i);
     }
   }
-  // Handle `struct_extract` and `struct_element_addr` instructions.
+  // Handle element projection instructions:
+  //   `struct_extract`, `struct_element_addr`, `ref_element_addr`.
   // - If the field is marked `@noDerivative`, do not set the result as
   //   varied because it does not need a derivative.
   // - Otherwise, propagate variedness from operand to result as usual.
-#define PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(INST)                           \
-  else if (auto *sei = dyn_cast<INST##Inst>(inst)) {                           \
-    if (isVaried(sei->getOperand(), i) &&                                      \
-        !sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())         \
-      setVariedAndPropagateToUsers(sei, i);                                    \
+#define PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(INST)                          \
+  else if (auto *projInst = dyn_cast<INST##Inst>(inst)) {                      \
+    if (isVaried(projInst->getOperand(), i) &&                                 \
+        !projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())    \
+      setVariedAndPropagateToUsers(projInst, i);                               \
   }
-  PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(StructExtract)
-  PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(StructElementAddr)
+  PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(StructExtract)
+  PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(StructElementAddr)
+  PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(RefElementAddr)
 #undef PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION
   // Handle `br`.
   else if (auto *bi = dyn_cast<BranchInst>(inst)) {
@@ -222,13 +224,14 @@ static Optional<AccessorKind> getAccessorKind(SILFunction *fn) {
 void DifferentiableActivityInfo::propagateVariedInwardsThroughProjections(
     SILValue value, unsigned independentVariableIndex) {
   auto i = independentVariableIndex;
-  // Skip `@noDerivative` struct projections.
+  // Skip `@noDerivative` projections.
 #define SKIP_NODERIVATIVE(INST)                                                \
-  if (auto *sei = dyn_cast<INST##Inst>(value))                                 \
-    if (sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())          \
+  if (auto *projInst = dyn_cast<INST##Inst>(value))                            \
+    if (projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())     \
       return;
   SKIP_NODERIVATIVE(StructExtract)
   SKIP_NODERIVATIVE(StructElementAddr)
+  SKIP_NODERIVATIVE(RefElementAddr)
 #undef SKIP_NODERIVATIVE
   // Set value as varied and propagate to users.
   setVariedAndPropagateToUsers(value, i);
@@ -274,7 +277,8 @@ void DifferentiableActivityInfo::setUsefulAndPropagateToOperands(
   // Skip already-useful values to prevent infinite recursion.
   if (isUseful(value, dependentVariableIndex))
     return;
-  if (value->getType().isAddress()) {
+  if (value->getType().isAddress() ||
+      value->getType().getClassOrBoundGenericClass()) {
     propagateUsefulThroughAddress(value, dependentVariableIndex);
     return;
   }
@@ -331,15 +335,16 @@ void DifferentiableActivityInfo::propagateUseful(
   PROPAGATE_USEFUL_THROUGH_STORE(CopyAddr)
   PROPAGATE_USEFUL_THROUGH_STORE(UnconditionalCheckedCastAddr)
 #undef PROPAGATE_USEFUL_THROUGH_STORE
-  // Handle struct element extraction, skipping `@noDerivative` fields:
-  //   `struct_extract`, `struct_element_addr`.
-#define PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(INST)                       \
-  else if (auto *sei = dyn_cast<INST##Inst>(inst)) {                           \
-    if (!sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())         \
-      setUsefulAndPropagateToOperands(sei->getOperand(), i);                   \
+  // Handle element projections, skipping `@noDerivative` fields:
+  //   `struct_extract`, `struct_element_addr`, `ref_element_addr`.
+#define PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(INST)                      \
+  else if (auto *projInst = dyn_cast<INST##Inst>(inst)) {                      \
+    if (!projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())    \
+      setUsefulAndPropagateToOperands(projInst->getOperand(), i);              \
   }
-  PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(StructExtract)
-  PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(StructElementAddr)
+  PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(StructExtract)
+  PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(StructElementAddr)
+  PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(RefElementAddr)
 #undef PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION
   // Handle everything else.
   else {
@@ -350,7 +355,8 @@ void DifferentiableActivityInfo::propagateUseful(
 
 void DifferentiableActivityInfo::propagateUsefulThroughAddress(
     SILValue value, unsigned dependentVariableIndex) {
-  assert(value->getType().isAddress());
+  assert(value->getType().isAddress() ||
+         value->getType().getClassOrBoundGenericClass());
   // Skip already-useful values to prevent infinite recursion.
   if (isUseful(value, dependentVariableIndex))
     return;
@@ -364,13 +370,15 @@ void DifferentiableActivityInfo::propagateUsefulThroughAddress(
     propagateUseful(use->getUser(), dependentVariableIndex);
     for (auto res : use->getUser()->getResults()) {
 #define SKIP_NODERIVATIVE(INST)                                                \
-  if (auto *sei = dyn_cast<INST##Inst>(res))                                   \
-    if (sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())          \
+  if (auto *projInst = dyn_cast<INST##Inst>(res))                              \
+    if (projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>())     \
       continue;
       SKIP_NODERIVATIVE(StructExtract)
       SKIP_NODERIVATIVE(StructElementAddr)
+      SKIP_NODERIVATIVE(RefElementAddr)
 #undef SKIP_NODERIVATIVE
-      if (Projection::isAddressProjection(res) || isa<BeginAccessInst>(res))
+      if (Projection::isAddressProjection(res) || isa<BeginAccessInst>(res) ||
+          isa<BeginBorrowInst>(res))
         propagateUsefulThroughAddress(res, dependentVariableIndex);
     }
   }