[AutoDiff] Support ref_element_addr differentiation.

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

Teach activity analysis about `ref_element_addr`, handling it 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
```

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,10 @@ 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 differentiate value with loadable type %0 but address-only "
+     "'TangentVector' type %1", (Type, Type))
 NOTE(autodiff_enums_unsupported,none,
      "differentiating enum values is not yet supported", ())
 NOTE(autodiff_global_let_closure_not_differentiable,none,
@@ -543,9 +547,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

@@ -354,6 +354,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 +428,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,13 @@ 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 +305,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);
     }
   }

lib/SILOptimizer/Utils/Differentiation/PullbackEmitter.cpp

@@ -328,6 +328,7 @@ SILValue PullbackEmitter::getAdjointProjection(SILBasicBlock *origBB,
     auto adjSource = getAdjointBuffer(origBB, seai->getOperand());
     auto *tangentVectorDecl =
         adjSource->getType().getStructOrBoundGenericStruct();
+    // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
     auto tanFieldLookup =
         tangentVectorDecl->lookupDirect(seai->getField()->getName());
     assert(tanFieldLookup.size() == 1);
@@ -505,22 +506,52 @@ bool PullbackEmitter::run() {
       auto &domBBActiveValues = activeValues[domNode->getBlock()];
       bbActiveValues.append(domBBActiveValues.begin(), domBBActiveValues.end());
     }
+    // Booleans tracking whether active-value-related errors have been emitted.
+    // This prevents duplicate diagnostics for the same active values.
     bool diagnosedActiveEnumValue = false;
+    bool diagnosedActiveValueTangentValueCategoryIncompatible = false;
     // Mark the activity of a value if it has not yet been visited.
     auto markValueActivity = [&](SILValue v) {
       if (visited.count(v))
         return;
       visited.insert(v);
+      auto type = v->getType();
       // Diagnose active enum values. Differentiation of enum values requires
       // special adjoint value handling and is not yet supported. Diagnose
       // only the first active enum value to prevent too many diagnostics.
-      if (!diagnosedActiveEnumValue &&
-          v->getType().getEnumOrBoundGenericEnum()) {
+      if (!diagnosedActiveEnumValue && type.getEnumOrBoundGenericEnum()) {
         getContext().emitNondifferentiabilityError(
             v, getInvoker(), diag::autodiff_enums_unsupported);
         errorOccurred = true;
         diagnosedActiveEnumValue = true;
       }
+      // Diagnose active values whose value category is incompatible with their
+      // tangent types's value category.
+      //
+      // Let $L be a loadable type and $*A be an address-only type.
+      // Table of supported combinations:
+      //
+      // Original type | Tangent type | Currently supported?
+      // --------------|--------------|---------------------
+      // $L            | $L           | Yes (no mismatch)
+      // $*A           | $L           | Yes (can create $*L adjoint buffer)
+      // $L            | $*A          | No (cannot create $A adjoint value)
+      // $*A           | $*A          | Yes (no mismatch)
+      if (!diagnosedActiveValueTangentValueCategoryIncompatible) {
+        if (auto tanSpace = getTangentSpace(remapType(type).getASTType())) {
+          auto tanASTType = tanSpace->getCanonicalType();
+          auto &origTL = getTypeLowering(type.getASTType());
+          auto &tanTL = getTypeLowering(tanASTType);
+          if (!origTL.isAddressOnly() && tanTL.isAddressOnly()) {
+            getContext().emitNondifferentiabilityError(
+                v, getInvoker(),
+                diag::autodiff_loadable_value_addressonly_tangent_unsupported,
+                type.getASTType(), tanASTType);
+            diagnosedActiveValueTangentValueCategoryIncompatible = true;
+            errorOccurred = true;
+          }
+        }
+      }
       // Skip address projections.
       // Address projections do not need their own adjoint buffers; they
       // become projections into their adjoint base buffer.
@@ -1050,7 +1081,7 @@ PullbackEmitter::getArrayAdjointElementBuffer(SILValue arrayAdjoint,
   auto subMap = SubstitutionMap::get(subscriptFnGenSig, {eltTanType},
                                      {addArithConf, diffConf});
   // %elt_adj = alloc_stack $T.TangentVector
-  // Create a local allocation and register it.
+  // Create and register a local allocation.
   localAllocBuilder.setInsertionPoint(
       getPullback().getEntryBlock(),
       getNextFunctionLocalAllocationInsertionPoint());
@@ -1293,6 +1324,7 @@ void PullbackEmitter::visitStructExtractInst(StructExtractInst *sei) {
   assert(!getTypeLowering(tangentVectorTy).isAddressOnly());
   auto tangentVectorSILTy = SILType::getPrimitiveObjectType(tangentVectorTy);
   auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
+  // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
   assert(tangentVectorDecl);
   // Find the corresponding field in the tangent space.
   VarDecl *tanField = nullptr;
@@ -1343,6 +1375,53 @@ void PullbackEmitter::visitStructExtractInst(StructExtractInst *sei) {
   }
 }
 
+void PullbackEmitter::visitRefElementAddrInst(RefElementAddrInst *reai) {
+  auto *bb = reai->getParent();
+  auto adjBuf = getAdjointBuffer(bb, reai);
+  auto classTy = remapType(reai->getOperand()->getType()).getASTType();
+  auto tangentVectorTy =
+      getTangentSpace(classTy)->getType()->getCanonicalType();
+  assert(!getTypeLowering(tangentVectorTy).isAddressOnly());
+  auto tangentVectorSILTy = SILType::getPrimitiveObjectType(tangentVectorTy);
+  auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
+  // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
+  assert(tangentVectorDecl);
+  // Look up the corresponding field in the tangent space by name.
+  VarDecl *tanField = nullptr;
+  auto tanFieldLookup =
+      tangentVectorDecl->lookupDirect(reai->getField()->getName());
+  if (tanFieldLookup.empty()) {
+    getContext().emitNondifferentiabilityError(
+        reai, getInvoker(),
+        diag::autodiff_stored_property_no_corresponding_tangent,
+        reai->getClassDecl()->getNameStr(), reai->getField()->getNameStr());
+    errorOccurred = true;
+    return;
+  }
+  tanField = cast<VarDecl>(tanFieldLookup.front());
+  // Accumulate adjoint for the `ref_element_addr` operand.
+  SmallVector<AdjointValue, 8> eltVals;
+  for (auto *field : tangentVectorDecl->getStoredProperties()) {
+    if (field == tanField) {
+      auto adjElt = builder.emitLoadValueOperation(
+          reai->getLoc(), adjBuf, LoadOwnershipQualifier::Copy);
+      eltVals.push_back(makeConcreteAdjointValue(adjElt));
+      recordTemporary(adjElt);
+    } else {
+      auto substMap = tangentVectorTy->getMemberSubstitutionMap(
+          field->getModuleContext(), field);
+      auto fieldTy = field->getType().subst(substMap);
+      auto fieldSILTy = getContext().getTypeConverter().getLoweredType(
+          fieldTy, TypeExpansionContext::minimal());
+      assert(fieldSILTy.isObject());
+      eltVals.push_back(makeZeroAdjointValue(fieldSILTy));
+    }
+  }
+  addAdjointValue(bb, reai->getOperand(),
+                  makeAggregateAdjointValue(tangentVectorSILTy, eltVals),
+                  reai->getLoc());
+}
+
 void PullbackEmitter::visitTupleInst(TupleInst *ti) {
   auto *bb = ti->getParent();
   auto av = getAdjointValue(bb, ti);
@@ -1558,7 +1637,6 @@ void PullbackEmitter::visitUnconditionalCheckedCastAddrInst(
     errorOccurred = true;                                                      \
     return;                                                                    \
   }
-NOT_DIFFERENTIABLE(RefElementAddr, autodiff_class_property_not_supported)
 #undef NOT_DIFFERENTIABLE
 
 AdjointValue PullbackEmitter::makeZeroAdjointValue(SILType type) {

test/AutoDiff/downstream/activity_analysis.swift

@@ -569,6 +569,35 @@ func testBeginApplyActiveButInitiallyNonactiveInoutArgument(x: Float) -> Float {
 // CHECK: [NONE]   %34 = apply %33<Float>(%32, %29, %31) : $@convention(method) <τ_0_0> (Int, @guaranteed Array<τ_0_0>) -> @out τ_0_0
 // CHECK: [ACTIVE]   %35 = load [trivial] %32 : $*Float
 
+//===----------------------------------------------------------------------===//
+// Class differentiation
+//===----------------------------------------------------------------------===//
+
+class C: Differentiable {
+  @differentiable
+  var float: Float
+
+  init(_ float: Float) {
+    self.float = float
+  }
+
+  @differentiable
+  func method(_ x: Float) -> Float {
+    x * float
+  }
+
+// CHECK-LABEL: [AD] Activity info for ${{.*}}1CC6methodyS2fF at (source=0 parameters=(0 1))
+// CHECK: bb0:
+// CHECK: [ACTIVE] %0 = argument of bb0 : $Float
+// CHECK: [ACTIVE] %1 = argument of bb0 : $C
+// CHECK: [USEFUL]   %4 = metatype $@thin Float.Type
+// CHECK: [VARIED]   %5 = class_method %1 : $C, #C.float!getter.1 : (C) -> () -> Float, $@convention(method) (@guaranteed C) -> Float
+// CHECK: [ACTIVE]   %6 = apply %5(%1) : $@convention(method) (@guaranteed C) -> Float
+// CHECK: [NONE]   // function_ref static Float.* infix(_:_:)
+// CHECK:   %7 = function_ref @$sSf1moiyS2f_SftFZ : $@convention(method) (Float, Float, @thin Float.Type) -> Float
+// CHECK: [ACTIVE]   %8 = apply %7(%0, %6, %4) : $@convention(method) (Float, Float, @thin Float.Type) -> Float
+}
+
 //===----------------------------------------------------------------------===//
 // Enum differentiation
 //===----------------------------------------------------------------------===//