Fix incorrect ref_element_addr derivative values.

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.

Fixes incorrect `ref_element_addr` zero derivatives.
Add class initializer differentiation tests.

Author: dan-zheng

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);
 

lib/SILOptimizer/Utils/Differentiation/PullbackEmitter.cpp

@@ -19,6 +19,7 @@
 #define DEBUG_TYPE "differentiation"
 
 #include "swift/SILOptimizer/Utils/Differentiation/PullbackEmitter.h"
+#include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SILOptimizer/PassManager/PrettyStackTrace.h"
 #include "swift/SILOptimizer/Utils/Differentiation/ADContext.h"
@@ -351,6 +352,43 @@ SILValue PullbackEmitter::getAdjointProjection(SILBasicBlock *origBB,
     }
     return builder.createTupleElementAddr(teai->getLoc(), adjSource, adjIndex);
   }
+  // Handle `ref_element_addr`.
+  // Adjoint projection: a local allocation initialized with the corresponding
+  // field value from the class's base adjoint value.
+  if (auto *reai = dyn_cast<RefElementAddrInst>(originalProjection)) {
+    assert(!reai->getField()->getAttrs().hasAttribute<NoDerivativeAttr>() &&
+           "`@noDerivative` class projections should never be active");
+    auto loc = reai->getLoc();
+    // Get the class operand, stripping `begin_borrow`.
+    auto classOperand = stripBorrow(reai->getOperand());
+    // Get the class operand's adjoint value. Currently, it must be a
+    // `TangentVector` struct.
+    auto adjClass =
+        materializeAdjointDirect(getAdjointValue(origBB, classOperand), loc);
+    auto *tangentVectorDecl =
+        adjClass->getType().getStructOrBoundGenericStruct();
+    // TODO(TF-970): Replace assertions below with diagnostics.
+    assert(tangentVectorDecl && "`TangentVector` of a class must be a struct");
+    auto tanFieldLookup =
+        tangentVectorDecl->lookupDirect(reai->getField()->getName());
+    assert(tanFieldLookup.size() == 1 &&
+           "Class `TangentVector` must have field of the same name");
+    auto *tanField = cast<VarDecl>(tanFieldLookup.front());
+    // Create a local allocation for the element adjoint buffer.
+    auto eltTanType = tanField->getValueInterfaceType()->getCanonicalType();
+    auto eltTanSILType = SILType::getPrimitiveAddressType(eltTanType);
+    auto *eltAdjBuffer = createFunctionLocalAllocation(eltTanSILType, loc);
+    builder.emitScopedBorrowOperation(
+        loc, adjClass, [&](SILValue borrowedAdjClass) {
+          // Initialize the element adjoint buffer with the base adjoint value.
+          auto *adjElt =
+              builder.createStructExtract(loc, borrowedAdjClass, tanField);
+          auto adjEltCopy = builder.emitCopyValueOperation(loc, adjElt);
+          builder.emitStoreValueOperation(loc, adjEltCopy, eltAdjBuffer,
+                                          StoreOwnershipQualifier::Init);
+        });
+    return eltAdjBuffer;
+  }
   // Handle `begin_access`.
   // Adjoint projection: the base adjoint buffer itself.
   if (auto *bai = dyn_cast<BeginAccessInst>(originalProjection)) {
@@ -412,6 +450,20 @@ PullbackEmitter::getNextFunctionLocalAllocationInsertionPoint() {
   return lastLocalAlloc->getDefiningInstruction()->getIterator();
 }
 
+AllocStackInst *
+PullbackEmitter::createFunctionLocalAllocation(SILType type, SILLocation loc) {
+  // Set insertion point for local allocation builder: before the last local
+  // allocation, or at the start of the pullback function's entry if no local
+  // allocations exist yet.
+  localAllocBuilder.setInsertionPoint(
+      getPullback().getEntryBlock(),
+      getNextFunctionLocalAllocationInsertionPoint());
+  // Create and return local allocation.
+  auto *alloc = localAllocBuilder.createAllocStack(loc, type);
+  functionLocalAllocations.push_back(alloc);
+  return alloc;
+}
+
 SILValue &PullbackEmitter::getAdjointBuffer(SILBasicBlock *origBB,
                                             SILValue originalBuffer) {
   assert(originalBuffer->getType().isAddress());
@@ -425,25 +477,19 @@ SILValue &PullbackEmitter::getAdjointBuffer(SILBasicBlock *origBB,
   if (auto adjProj = getAdjointProjection(origBB, originalBuffer))
     return (bufferMap[{origBB, originalBuffer}] = adjProj);
 
+  auto bufObjectType = getRemappedTangentType(originalBuffer->getType());
   // Set insertion point for local allocation builder: before the last local
   // allocation, or at the start of the pullback function's entry if no local
   // allocations exist yet.
-  localAllocBuilder.setInsertionPoint(
-      getPullback().getEntryBlock(),
-      getNextFunctionLocalAllocationInsertionPoint());
-  // Allocate local buffer and initialize to zero.
-  auto bufObjectType = getRemappedTangentType(originalBuffer->getType());
-  auto *newBuf = localAllocBuilder.createAllocStack(
-      RegularLocation::getAutoGeneratedLocation(), bufObjectType);
+  auto *newBuf = createFunctionLocalAllocation(
+      bufObjectType, RegularLocation::getAutoGeneratedLocation());
   // Temporarily change global builder insertion point and emit zero into the
-  // local buffer.
+  // local allocation.
   auto insertionPoint = builder.getInsertionBB();
   builder.setInsertionPoint(localAllocBuilder.getInsertionBB(),
                             localAllocBuilder.getInsertionPoint());
   emitZeroIndirect(bufObjectType.getASTType(), newBuf, newBuf->getLoc());
   builder.setInsertionPoint(insertionPoint);
-  // Register the local buffer.
-  functionLocalAllocations.push_back(newBuf);
   return (insertion.first->getSecond() = newBuf);
 }
 
@@ -1082,13 +1128,9 @@ PullbackEmitter::getArrayAdjointElementBuffer(SILValue arrayAdjoint,
                                      {addArithConf, diffConf});
   // %elt_adj = alloc_stack $T.TangentVector
   // Create and register a local allocation.
-  localAllocBuilder.setInsertionPoint(
-      getPullback().getEntryBlock(),
-      getNextFunctionLocalAllocationInsertionPoint());
-  auto *eltAdjBuffer = localAllocBuilder.createAllocStack(loc, eltTanSILType);
-  functionLocalAllocations.push_back(eltAdjBuffer);
+  auto *eltAdjBuffer = createFunctionLocalAllocation(eltTanSILType, loc);
   // Temporarily change global builder insertion point and emit zero into the
-  // local buffer.
+  // local allocation.
   auto insertionPoint = builder.getInsertionBB();
   builder.setInsertionPoint(localAllocBuilder.getInsertionBB(),
                             localAllocBuilder.getInsertionPoint());

test/AutoDiff/downstream/class_differentiation.swift

@@ -34,13 +34,17 @@ ClassTests.test("TrivialMember") {
     static func controlFlow(_ c1: C, _ c2: C, _ flag: Bool) -> Float {
       var result: Float = 0
       if flag {
-        result = c1.float * c2.float
+        var c3 = C(c1.float * c2.float)
+        result = c3.float
       } else {
         result = c2.float * c1.float
       }
       return result
     }
   }
+  // Test class initializer differentiation.
+  expectEqual(10, pullback(at: 3, in: { C($0) })(.init(float: 10)))
+  // Test class method differentiation.
   expectEqual((.init(float: 3), 10), gradient(at: C(10), 3, in: { c, x in c.method(x) }))
   expectEqual(.init(float: 0), gradient(at: C(10), in: { c in c.testNoDerivative() }))
   expectEqual((.init(float: 20), .init(float: 10)),
@@ -72,6 +76,9 @@ ClassTests.test("NontrivialMember") {
       return result
     }
   }
+  // Test class initializer differentiation.
+  expectEqual(10, pullback(at: 3, in: { C($0) })(.init(float: 10)))
+  // Test class method differentiation.
   expectEqual((.init(float: 3), 10), gradient(at: C(10), 3, in: { c, x in c.method(x) }))
   expectEqual((.init(float: 20), .init(float: 10)),
               gradient(at: C(10), C(20), in: { c1, c2 in C.controlFlow(c1, c2, true) }))
@@ -94,6 +101,9 @@ ClassTests.test("AddressOnlyTangentVector") {
       stored
     }
   }
+  // Test class initializer differentiation.
+  expectEqual(10, pullback(at: 3, in: { C<Float>($0) })(.init(float: 10)))
+  // Test class method differentiation.
   expectEqual((.init(stored: Float(3)), 10),
               gradient(at: C<Float>(3), 3, in: { c, x in c.method(x) }))
 }