[AutoDiff] Start supporting loadable types with address-only tangents. (#32540)

Previously, PullbackEmitter assumed that original values' value category
matches their `TangentVector` types' value category. This was problematic
for loadable types with address-only `TangentVector` types.

Now, PullbackEmitter starts to support differentiation of loadable types with
address-only `TangentVector` types. This patch focuses on supporting and testing
class types, more support can be added incrementally.

Resolves TF-1149.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -496,12 +496,6 @@ NOTE(autodiff_cannot_differentiate_through_multiple_results,none,
      "cannot differentiate through multiple results", ())
 NOTE(autodiff_cannot_differentiate_through_inout_arguments,none,
      "cannot differentiate through 'inout' arguments", ())
-// 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_stored_property_parent_not_differentiable,none,

include/swift/SILOptimizer/Differentiation/PullbackEmitter.h

@@ -218,6 +218,9 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
 
   Optional<TangentSpace> getTangentSpace(CanType type);
 
+  /// Returns the tangent value category of the given value.
+  SILValueCategory getTangentValueCategory(SILValue v);
+
   /// Assuming the given type conforms to `Differentiable` after remapping,
   /// returns the associated tangent space type.
   SILType getRemappedTangentType(SILType type);
@@ -264,6 +267,8 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   SILValue getAdjointProjection(SILBasicBlock *origBB,
                                 SILValue originalProjection);
 
+  SILValue &getAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer);
+
   SILBasicBlock::iterator getNextFunctionLocalAllocationInsertionPoint();
 
   /// Creates and returns a local allocation with the given type.
@@ -273,8 +278,6 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   /// `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`.
   void addToAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer,

lib/SILOptimizer/Differentiation/PullbackEmitter.cpp

@@ -380,6 +380,23 @@ SILFunction *getOrCreateReabstractionThunk(SILOptFunctionBuilder &fb,
   for (unsigned resIdx : range(toType->getNumResults())) {
     auto fromRes = fromConv.getResults()[resIdx];
     auto toRes = toConv.getResults()[resIdx];
+    // Check function-typed results.
+    if (isa<SILFunctionType>(fromRes.getInterfaceType()) &&
+        isa<SILFunctionType>(toRes.getInterfaceType())) {
+      auto fromFnType = cast<SILFunctionType>(fromRes.getInterfaceType());
+      auto toFnType = cast<SILFunctionType>(toRes.getInterfaceType());
+      auto fromUnsubstFnType = fromFnType->getUnsubstitutedType(module);
+      auto toUnsubstFnType = toFnType->getUnsubstitutedType(module);
+      // If unsubstituted function types are not equal, perform reabstraction.
+      if (fromUnsubstFnType != toUnsubstFnType) {
+        auto fromFn = *fromDirResultsIter++;
+        auto newFromFn = reabstractFunction(
+            builder, fb, loc, fromFn, toFnType,
+            [](SubstitutionMap substMap) { return substMap; });
+        results.push_back(newFromFn);
+        continue;
+      }
+    }
     // No abstraction mismatch.
     if (fromRes.isFormalIndirect() == toRes.isFormalIndirect()) {
       // If result types are direct, add call result as direct thunk result.

lib/SILOptimizer/Differentiation/Thunk.cpp

@@ -1170,6 +1170,23 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
           loc, derivativeFn,
           SILType::getPrimitiveObjectType(expectedDerivativeFnTy));
     }
+    // If derivative function value's type is not ABI-compatible with the
+    // expected derivative function type (i.e. parameter and result conventions
+    // do not match), perform reabstraction.
+    auto abiCompatibility = expectedDerivativeFnTy->isABICompatibleWith(
+        derivativeFn->getType().castTo<SILFunctionType>(), *dfi->getFunction());
+    if (!abiCompatibility.isCompatible()) {
+      SILOptFunctionBuilder fb(context.getTransform());
+      auto newDerivativeFn = reabstractFunction(
+          builder, fb, loc, derivativeFn, expectedDerivativeFnTy,
+          [](SubstitutionMap substMap) { return substMap; });
+      derivativeFn = newDerivativeFn;
+      assert(expectedDerivativeFnTy
+                 ->isABICompatibleWith(
+                     derivativeFn->getType().castTo<SILFunctionType>(),
+                     *dfi->getFunction())
+                 .isCompatible());
+    }
 
     derivativeFns.push_back(derivativeFn);
   }

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -136,20 +136,14 @@ func testMultipleDiffAttrsClass<C: ClassMethodMultipleDifferentiableAttribute>(
 
 // TF-1149: Test class with loadable type but address-only `TangentVector` type.
 class C<T: Differentiable>: Differentiable {
-  // expected-error @+1 {{function is not differentiable}}
   @differentiable
-  // expected-note @+2 {{when differentiating this function definition}}
-  // expected-note @+1 {{cannot yet differentiate value whose type 'C<T>' has a compile-time known size, but whose 'TangentVector' contains stored properties of unknown size; consider modifying 'C<τ_0_0>.TangentVector' to use fewer generic parameters in stored properties}}
   var stored: T
 
   init(_ stored: T) {
     self.stored = stored
   }
 
-  // expected-error @+1 {{function is not differentiable}}
   @differentiable
-  // expected-note @+2 {{when differentiating this function definition}}
-  // expected-note @+1 {{cannot yet differentiate value whose type 'C<T>' has a compile-time known size, but whose 'TangentVector' contains stored properties of unknown size; consider modifying 'C<τ_0_0>.TangentVector' to use fewer generic parameters in stored properties}}
   func method(_ x: T) -> T {
     stored
   }

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -0,0 +1,67 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+import DifferentiationUnittest
+
+var AddressOnlyTangentVectorTests = TestSuite("AddressOnlyTangentVector")
+
+// TF-1149: Test loadable class type with an address-only `TangentVector` type.
+
+AddressOnlyTangentVectorTests.test("LoadableClassAddressOnlyTangentVector") {
+  final class LoadableClass<T: Differentiable>: Differentiable {
+    @differentiable
+    var stored: T
+
+    @differentiable
+    init(_ stored: T) {
+      self.stored = stored
+    }
+
+    @differentiable
+    func method(_ x: T) -> T {
+      stored
+    }
+  }
+
+  @differentiable
+  func projection<T: Differentiable>(_ s: LoadableClass<T>) -> T {
+    var x = s.stored
+    return x
+  }
+  expectEqual(.init(stored: 1), gradient(at: LoadableClass<Float>(10), in: projection))
+
+  @differentiable
+  func tuple<T: Differentiable>(_ s: LoadableClass<T>) -> T {
+    var tuple = (s, (s, s))
+    return tuple.1.0.stored
+  }
+  expectEqual(.init(stored: 1), gradient(at: LoadableClass<Float>(10), in: tuple))
+
+  @differentiable
+  func conditional<T: Differentiable>(_ s: LoadableClass<T>) -> T {
+    var tuple = (s, (s, s))
+    if false {}
+    return tuple.1.0.stored
+  }
+  expectEqual(.init(stored: 1), gradient(at: LoadableClass<Float>(10), in: conditional))
+
+  @differentiable
+  func loop<T: Differentiable>(_ array: [LoadableClass<T>]) -> T {
+    var result: [LoadableClass<T>] = []
+    for i in withoutDerivative(at: array.indices) {
+      result.append(array[i])
+    }
+    return result[0].stored
+  }
+  expectEqual([.init(stored: 1)], gradient(at: [LoadableClass<Float>(10)], in: loop))
+
+  @differentiable
+  func arrayLiteral<T: Differentiable>(_ s: LoadableClass<T>) -> T {
+    var result: [[LoadableClass<T>]] = [[s, s]]
+    return result[0][1].stored
+  }
+  expectEqual(.init(stored: 1), gradient(at: LoadableClass<Float>(10), in: arrayLiteral))
+}
+
+runAllTests()

test/AutoDiff/validation-test/address_only_tangentvector.swift

@@ -98,10 +98,7 @@ PropertyWrapperTests.test("GenericStruct") {
   */
 }
 
-// FIXME(TF-1149): Cannot differentiate active value with loadable type but
-// address-only tangent type. Triggered by marking properties with
-// `@differentiable`, which triggers derivative vtable thunk entries.
-/*
+// TF-1149: Test class with loadable type but address-only `TangentVector` type.
 class Class: Differentiable {
   @differentiable
   @Wrapper @Wrapper var x: Tracked<Float> = 10
@@ -124,8 +121,15 @@ PropertyWrapperTests.test("SimpleClass") {
     c.x = c.x * x * c.z
     return c.x
   }
+  // FIXME(TF-1175): Class operands should always be marked active.
+  // This is relevant for `Class.x.setter`, which has type
+  // `$@convention(method) (@in Tracked<Float>, @guaranteed Class) -> ()`.
+  expectEqual((.init(x: 1, y: 0, z: 0), 0),
+              gradient(at: Class(), 2, in: setter))
+  /*
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: Class(), 2, in: setter))
+  */
 
   // TODO(SR-12640): Support `modify` accessors.
   /*
@@ -138,7 +142,6 @@ PropertyWrapperTests.test("SimpleClass") {
               gradient(at: Class(), 2, in: modify))
   */
 }
-*/
 
 // From: https://github.com/apple/swift-evolution/blob/master/proposals/0258-property-wrappers.md#proposed-solution
 // Tests the following functionality: