[AutoDiff] Support @differentiable class initializers.

Support `@differentiable` class initializers.
- Add special logic in `SILGenModule::getOrCreateCustomDerivativeThunk` to
  thunk user-defined class initializer derivative functions to the expected
  derivative type for class initializers.
- Make minor vtable SILGen changes for class initializer derivatives.

Also support differentiation for class-related casting instructions:
`unchecked_ref_cast` and `upcast`.

These instructions are generated when calling `super.init` or referencing
inherited superclass members.

Resolves SR-12151 and SR-12153.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -3004,10 +3004,6 @@ ERROR(differentiable_attr_stored_property_variable_unsupported,none,
 ERROR(differentiable_attr_class_member_no_dynamic_self,none,
       "'@differentiable' attribute cannot be declared on class methods "
       "returning 'Self'", ())
-// TODO(TF-654): Remove when differentiation supports class initializers.
-ERROR(differentiable_attr_class_init_not_yet_supported,none,
-      "'@differentiable' attribute does not yet support class initializers",
-      ())
 ERROR(differentiable_attr_empty_where_clause,none,
       "empty 'where' clause in '@differentiable' attribute", ())
 // SWIFT_ENABLE_TENSORFLOW

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

@@ -434,6 +434,16 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   void visitUnconditionalCheckedCastAddrInst(
       UnconditionalCheckedCastAddrInst *uccai);
 
+  /// Handle `unchecked_ref_cast` instruction.
+  ///   Original: y = unchecked_ref_cast x
+  ///    Adjoint: adj[x] += adj[y] (assuming x' and y' have the same type)
+  void visitUncheckedRefCastInst(UncheckedRefCastInst *urci);
+
+  /// Handle `upcast` instruction.
+  ///   Original: y = upcast x
+  ///    Adjoint: adj[x] += adj[y] (assuming x' and y' have the same type)
+  void visitUpcastInst(UpcastInst *ui);
+
 #define NOT_DIFFERENTIABLE(INST, DIAG) void visit##INST##Inst(INST##Inst *inst);
 #undef NOT_DIFFERENTIABLE
 

lib/AST/Type.cpp

@@ -792,6 +792,20 @@ Type TypeBase::replaceCovariantResultType(Type newResultType,
       return OptionalType::get(
           objectType->replaceCovariantResultType(newResultType, uncurryLevel));
     }
+    // SWIFT_ENABLE_TENSORFLOW
+    // Special logic to handle JVP/VJP derivative functions of `Self`-returning
+    // methods. This logic is hacky and not robust at all.
+    // Consider adding a `bool isAutoDiffDerivative` flag for triggering this
+    // logic, or creating a separate dedicated helper
+    // `TypeBase::replaceCovariantResultTypeForAutoDiffDerivative`.
+    if (auto tupleType = dyn_cast<TupleType>(this)) {
+      assert(tupleType->getNumElements() == 2 &&
+             "Tuple result is expected only for derivative functions, which "
+             "return a two-element tuple");
+      return TupleType::get({newResultType, tupleType->getElement(1)},
+                            getASTContext());
+    }
+    // SWIFT_ENABLE_TENSORFLOW END
 
     return newResultType;
   }

lib/SIL/SILFunctionType.cpp

@@ -2705,17 +2705,15 @@ TypeConverter::getConstantInfo(TypeExpansionContext expansion,
                                             loweredInterfaceType);
 
   // SWIFT_ENABLE_TENSORFLOW
-  // In the case of autodiff derivative functions, the above computations
-  // determine `silFnType` by first computing the derivative function type at
-  // the AST level and then lowering that. Unfortunately, the actual
-  // SILFunctionType for the function is determined by first lowering the
-  // function's AST type, and then computing the derivative function type at the
-  // SIL level. "Lowering" does not commute with "getting the autodiff
-  // associated type", so these two computations produce different results.
-  // Therefore `silFnType` is not the actual type of the function that
-  // `constant` refers to.
+  // For derivative functions, the above computations determine `silFnType`
+  // by first computing the derivative AST function type and then lowering it to
+  // SIL. Unfortunately, the expected derivative SIL function type is determined
+  // by first lowering the original function's AST type, and then computing its
+  // SIL derivative function type. "Lowering" does not commute with "getting the
+  // derivative type", so these two computations produce different results.
+  // Therefore, `silFnType` is not the expected SIL derivative function type.
   //
-  // We hackily fix this problem by redoing the computation in the right order.
+  // We fix this problem by performing the computation in the right order.
   if (auto *autoDiffFuncId = constant.autoDiffDerivativeFunctionIdentifier) {
     auto origFnConstantInfo = getConstantInfo(
         TypeExpansionContext::minimal(), constant.asAutoDiffOriginalFunction());
@@ -2725,6 +2723,7 @@ TypeConverter::getConstantInfo(TypeExpansionContext expansion,
         loweredIndices, /*resultIndex*/ 0, autoDiffFuncId->getKind(),
         *this, LookUpConformanceInModule(&M));
   }
+  // SWIFT_ENABLE_TENSORFLOW END
 
   LLVM_DEBUG(llvm::dbgs() << "lowering type for constant ";
              constant.print(llvm::dbgs());
@@ -2938,6 +2937,30 @@ TypeConverter::getConstantOverrideInfo(TypeExpansionContext context,
       *this, context, basePattern, bridgedTypes.Uncurried, base, derived,
       /*reqt subs*/ None, ProtocolConformanceRef());
 
+  // SWIFT_ENABLE_TENSORFLOW
+  // For derivative functions, the above computations determine `fnTy`
+  // by first computing the derivative AST function type and then lowering it to
+  // SIL. Unfortunately, the expected derivative SIL function type is determined
+  // by first lowering the original function's AST type, and then computing its
+  // SIL derivative function type. "Lowering" does not commute with "getting the
+  // derivative type", so these two computations produce different results.
+  // Therefore, `fnTy` is not the expected SIL derivative function type.
+  //
+  // We fix this problem by performing the computation in the right order.
+  if (auto *derivedDerivID = derived.autoDiffDerivativeFunctionIdentifier) {
+    auto *baseDerivID = base.autoDiffDerivativeFunctionIdentifier;
+    assert(baseDerivID);
+    auto &overrideInfo =
+        getConstantOverrideInfo(context, derived.asAutoDiffOriginalFunction(),
+                                base.asAutoDiffOriginalFunction());
+    auto loweredIndices = autodiff::getLoweredParameterIndices(
+        derivedDerivID->getParameterIndices(), overrideInterfaceTy);
+    fnTy = overrideInfo.SILFnType->getAutoDiffDerivativeFunctionType(
+        loweredIndices, /*resultIndex*/ 0, derivedDerivID->getKind(), *this,
+        LookUpConformanceInModule(&M));
+  }
+  // SWIFT_ENABLE_TENSORFLOW END
+
   // Build the SILConstantInfo and cache it.
   auto resultBuf = Context.Allocate(sizeof(SILConstantInfo),
                                     alignof(SILConstantInfo));

lib/SILGen/SILGenPoly.cpp

@@ -3749,8 +3749,9 @@ SILFunction *SILGenModule::getOrCreateCustomDerivativeThunk(
   auto *thunk = fb.getOrCreateFunction(
       loc, name, customDerivativeFn->getLinkage(), thunkFnTy, IsBare,
       IsNotTransparent, customDerivativeFn->isSerialized(),
-      customDerivativeFn->isDynamicallyReplaceable(), customDerivativeFn->getEntryCount(),
-      IsThunk, customDerivativeFn->getClassSubclassScope());
+      customDerivativeFn->isDynamicallyReplaceable(),
+      customDerivativeFn->getEntryCount(), IsThunk,
+      customDerivativeFn->getClassSubclassScope());
   thunk->setInlineStrategy(AlwaysInline);
   if (!thunk->empty())
     return thunk;
@@ -3762,15 +3763,39 @@ SILFunction *SILGenModule::getOrCreateCustomDerivativeThunk(
   thunkSGF.collectThunkParams(loc, params, &indirectResults);
 
   auto *fnRef = thunkSGF.B.createFunctionRef(loc, customDerivativeFn);
-  auto fnRefType =
-      fnRef->getType().castTo<SILFunctionType>();
+  auto fnRefType = fnRef->getType().castTo<SILFunctionType>();
 
   // Collect thunk arguments, converting ownership.
   SmallVector<SILValue, 8> arguments;
   for (auto *indRes : indirectResults)
     arguments.push_back(indRes);
-  forwardFunctionArguments(thunkSGF, loc, fnRefType, params,
-                           arguments);
+  forwardFunctionArguments(thunkSGF, loc, fnRefType, params, arguments);
+
+  // Special support for thunking class initializer derivatives.
+  //
+  // User-defined custom derivatives take a metatype as the last parameter:
+  // - `$(Param0, Param1, ..., @thick Class.Type) -> (...)`
+  // But class initializers take an allocated instance as the last parameter:
+  // - `$(Param0, Param1, ..., @owned Class) -> (...)`
+  //
+  // Adjust forwarded arguments:
+  // - Pop the last `@owned Class` argument.
+  // - Create a `@thick Class.Type` value and pass it as the last argument.
+  auto *origAFD =
+      cast<AbstractFunctionDecl>(originalFn->getDeclContext()->getAsDecl());
+  if (isa<ConstructorDecl>(origAFD) &&
+      SILDeclRef(origAFD, SILDeclRef::Kind::Initializer).mangle() ==
+          originalFn->getName()) {
+    auto classArgument = arguments.pop_back_val();
+    auto *classDecl = classArgument->getType().getClassOrBoundGenericClass();
+    assert(classDecl && "Expected last argument to have class type");
+    auto classMetatype = MetatypeType::get(
+        classDecl->getDeclaredInterfaceType(), MetatypeRepresentation::Thick);
+    auto canClassMetatype = classMetatype->getCanonicalType();
+    auto *metatype = thunkSGF.B.createMetatype(
+        loc, SILType::getPrimitiveObjectType(canClassMetatype));
+    arguments.push_back(metatype);
+  }
   // Apply function argument.
   auto apply = thunkSGF.emitApplyWithRethrow(
       loc, fnRef, /*substFnType*/ fnRef->getType(),

lib/SILOptimizer/Utils/Differentiation/PullbackEmitter.cpp

@@ -1672,6 +1672,26 @@ void PullbackEmitter::visitUnconditionalCheckedCastAddrInst(
   emitZeroIndirect(destType.getASTType(), adjDest, uccai->getLoc());
 }
 
+void PullbackEmitter::visitUncheckedRefCastInst(UncheckedRefCastInst *urci) {
+  auto *bb = urci->getParent();
+  assert(urci->getOperand()->getType().isObject());
+  assert(getRemappedTangentType(urci->getOperand()->getType()) ==
+             getRemappedTangentType(urci->getType()) &&
+         "Operand/result must have the same `TangentVector` type");
+  auto adj = getAdjointValue(bb, urci);
+  addAdjointValue(bb, urci->getOperand(), adj, urci->getLoc());
+}
+
+void PullbackEmitter::visitUpcastInst(UpcastInst *ui) {
+  auto *bb = ui->getParent();
+  assert(ui->getOperand()->getType().isObject());
+  assert(getRemappedTangentType(ui->getOperand()->getType()) ==
+             getRemappedTangentType(ui->getType()) &&
+         "Operand/result must have the same `TangentVector` type");
+  auto adj = getAdjointValue(bb, ui);
+  addAdjointValue(bb, ui->getOperand(), adj, ui->getLoc());
+}
+
 #define NOT_DIFFERENTIABLE(INST, DIAG)                                         \
   void PullbackEmitter::visit##INST##Inst(INST##Inst *inst) {                  \
     getContext().emitNondifferentiabilityError(inst, getInvoker(),             \

lib/Sema/TypeCheckAttr.cpp

@@ -3975,26 +3975,16 @@ llvm::Expected<IndexSubset *> DifferentiableAttributeTypeCheckRequest::evaluate(
   // Diagnose if original function is an invalid class member.
   if (isOriginalClassMember) {
     // Class methods returning dynamic `Self` are not supported.
-    // (For class methods, dynamic `Self` is supported only as the single
-    // result - tuple-returning JVPs/VJPs would not type-check.)
-    if (auto *originalFn = dyn_cast<FuncDecl>(original)) {
-      if (originalFn->hasDynamicSelfResult()) {
+    // For class methods, dynamic `Self` is supported only as the single
+    // result - tuple-returning JVPs/VJPs would not type-check.
+    if (auto *originalFD = dyn_cast<FuncDecl>(original)) {
+      if (originalFD->hasDynamicSelfResult()) {
         diags.diagnose(attr->getLocation(),
                        diag::differentiable_attr_class_member_no_dynamic_self);
         attr->setInvalid();
         return nullptr;
       }
     }
-
-    // TODO(TF-654): Class initializers are not yet supported.
-    // Extra JVP/VJP type calculation logic is necessary because classes have
-    // both allocators and initializers.
-    if (auto *initDecl = dyn_cast<ConstructorDecl>(original)) {
-      diags.diagnose(attr->getLocation(),
-                     diag::differentiable_attr_class_init_not_yet_supported);
-      attr->setInvalid();
-      return nullptr;
-    }
   }
 
   // Resolve the derivative generic signature.

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -1080,8 +1080,6 @@ class Super: Differentiable {
 
   var base: Float
 
-  // NOTE(TF-654): Class initializers are not yet supported.
-  // expected-error @+1 {{'@differentiable' attribute does not yet support class initializers}}
   @differentiable
   init(base: Float) {
     self.base = base

test/AutoDiff/downstream/class_differentiation.swift

@@ -92,6 +92,7 @@ ClassTests.test("AddressOnlyTangentVector") {
     @differentiable
     var stored: T
 
+    @differentiable
     init(_ stored: T) {
       self.stored = stored
     }