Merge pull request #30312 from apple/tensorflow-stage-wip

Fix AutoDiff for substituted SIL function types.

The TF-1196 master issue tracks remaining improvements and fixes.

Author: dan-zheng

include/swift/AST/Decl.h

@@ -4710,6 +4710,10 @@ class AbstractStorageDecl : public ValueDecl {
   /// To ensure an accessor is always returned, use getSynthesizedAccessor().
   AccessorDecl *getOpaqueAccessor(AccessorKind kind) const;
 
+  /// Collect all opaque accessors.
+  ArrayRef<AccessorDecl*>
+    getOpaqueAccessors(llvm::SmallVectorImpl<AccessorDecl*> &scratch) const;
+
   /// Return an accessor that was written in source. Returns null if the
   /// accessor was not explicitly defined by the user.
   AccessorDecl *getParsedAccessor(AccessorKind kind) const;

include/swift/AST/DiagnosticsParse.def

@@ -824,7 +824,7 @@ ERROR(sil_box_expected_r_angle,none,
 ERROR(sil_function_subst_expected_l_angle,none,
       "expected '<' to begin SIL function type substitution list after 'for'", ())
 ERROR(sil_function_subst_expected_r_angle,none,
-      "expected '>' to begin SIL function type substitution list after 'for'", ())
+      "expected '>' to end SIL function type substitution list after 'for <...'", ())
 
 // Opaque types
 ERROR(opaque_mid_composition,none,

include/swift/AST/Type.h

@@ -47,6 +47,7 @@ class NormalProtocolConformance;
 class ProtocolConformanceRef;
 class ProtocolDecl;
 class ProtocolType;
+class SILModule;
 class StructDecl;
 class SubstitutableType;
 class SubstitutionMap;

include/swift/AST/Types.h

@@ -899,14 +899,26 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   /// Visit this type and the argument type in parallel, invoking the callback
   /// function with each archetype-to-substituted-type binding. The callback
   /// may return a new type to substitute into the result type, or return
-  /// CanType() to error out of the operation.
+  /// CanType() to error out of the operation. Each invocation of the callback
+  /// receives three arguments:
+  /// - The `orig` archetype from a position in `this` type.
+  /// - The `subst` type in the same structural position of `ty` that is trying to be bound
+  ///  to `orig`.
+  /// - The `upperBound` archetype, which if set, indicates the minimum set of constraints
+  ///  that any type substituted in this structural position must conform to. May be null,
+  ///  indicating an unconstrained context.
+  /// - If `upperBound` is set, then the `substConformances` array will contain the
+  ///  protocol conformances for `subst` to each of the protocol requirements
+  ///  on `upperBound` in `getConformsTo` order.
   ///
   /// Returns the substituted type, or a null CanType() if this type
   /// is not bindable to the substituted type, or the callback returns
   /// CanType().
   CanType substituteBindingsTo(Type ty,
-         llvm::function_ref<CanType(ArchetypeType*, CanType)> substFn);
-
+         llvm::function_ref<CanType(ArchetypeType *orig,
+                CanType subst,
+                ArchetypeType *upperBound,
+                ArrayRef<ProtocolConformanceRef> substConformances)> substFn);
 
   /// Determines whether this type is similar to \p other as defined by
   /// \p matchOptions.
@@ -1168,6 +1180,8 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   /// object type.
   TypeTraitResult canBeClass();
 
+  Type replaceSubstitutedSILFunctionTypesWithUnsubstituted(SILModule &M) const; // in SILType.cpp
+  
   /// Return the tangent space of the given type, if it exists. Otherwise,
   /// return `None`.
   Optional<TangentSpace>
@@ -3977,7 +3991,7 @@ class SILResultInfo {
   }
 
   ValueOwnershipKind
-  getOwnershipKind(SILFunction &) const; // in SILType.cpp
+  getOwnershipKind(SILFunction &, CanSILFunctionType fTy) const; // in SILType.cpp
 
   bool operator==(SILResultInfo rhs) const {
     return TypeAndConvention == rhs.TypeAndConvention;

include/swift/Basic/Algorithm.h

@@ -17,13 +17,12 @@
 #ifndef SWIFT_BASIC_FLAGGEDPOINTER_H
 #define SWIFT_BASIC_FLAGGEDPOINTER_H
 
+#include <algorithm>
 #include <cassert>
 
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
-#include "Algorithm.h"
-
 namespace swift {
 
 /// This class implements a pair of a pointer and boolean flag.
@@ -170,7 +169,7 @@ struct llvm::PointerLikeTypeTraits<
   enum {
     NumLowBitsAvailable = (BitPosition >= PtrTraits::NumLowBitsAvailable)
       ? PtrTraits::NumLowBitsAvailable
-      : (swift::min(int(BitPosition + 1),
+      : (std::min(int(BitPosition + 1),
         int(PtrTraits::NumLowBitsAvailable)) - 1)
   };
 };

include/swift/Basic/FlaggedPointer.h

@@ -194,7 +194,7 @@ namespace swift {
     bool BuildRequestDependencyGraph = false;
 
     /// Enable SIL type lowering
-    bool EnableSubstSILFunctionTypesForFunctionValues = false;
+    bool EnableSubstSILFunctionTypesForFunctionValues = true;
 
     /// Whether to diagnose an ephemeral to non-ephemeral conversion as an
     /// error.

include/swift/Basic/LangOptions.h

@@ -34,7 +34,6 @@
 #ifndef SWIFT_BASIC_PREFIXMAP_H
 #define SWIFT_BASIC_PREFIXMAP_H
 
-#include "swift/Basic/Algorithm.h"
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/type_traits.h"
@@ -53,8 +52,8 @@ template <class KeyElementType> class PrefixMapKeyPrinter;
 /// A map whose keys are sequences of comparable values, optimized for
 /// finding a mapped value for the longest matching initial subsequence.
 template <class KeyElementType, class ValueType,
-          size_t InlineKeyCapacity
-             = max<size_t>((sizeof(void*) - 1) / sizeof(KeyElementType), 1)>
+          size_t InlineKeyCapacity = std::max(
+              (sizeof(void *) - 1) / sizeof(KeyElementType), size_t(1))>
 class PrefixMap {
 public:
   using KeyType = ArrayRef<KeyElementType>;

include/swift/Basic/PrefixMap.h

@@ -195,6 +195,10 @@ class ApplySite {
   SILType getSubstCalleeSILType() const {
     FOREACH_IMPL_RETURN(getSubstCalleeSILType());
   }
+  void setSubstCalleeType(CanSILFunctionType t) {
+    FOREACH_IMPL_RETURN(setSubstCalleeType(t));
+  }
+
   /// Get the conventions of the callee with the applied substitutions.
   SILFunctionConventions getSubstCalleeConv() const {
     return SILFunctionConventions(getSubstCalleeType(), getModule());

include/swift/SIL/ApplySite.h

@@ -2859,10 +2859,13 @@ template<typename ImplClass>
 void SILCloner<ImplClass>::
 visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  Optional<SILType> explicitExtracteeType = None;
+  if (Inst->hasExplicitExtracteeType())
+    explicitExtracteeType = Inst->getType();
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiableFunctionExtract(
                 getOpLocation(Inst->getLoc()), Inst->getExtractee(),
-                getOpValue(Inst->getFunctionOperand())));
+                getOpValue(Inst->getFunctionOperand()), explicitExtracteeType));
 }
 
 template<typename ImplClass>

include/swift/SIL/SILCloner.h

@@ -1933,6 +1933,11 @@ class ApplyInstBase<Impl, Base, false> : public Base {
   SILType getSubstCalleeSILType() const {
     return SubstCalleeType;
   }
+  
+  void setSubstCalleeType(CanSILFunctionType t) {
+    SubstCalleeType = SILType::getPrimitiveObjectType(t);
+  }
+  
   SILFunctionConventions getSubstCalleeConv() const {
     return SILFunctionConventions(getSubstCalleeType(), this->getModule());
   }
@@ -4352,6 +4357,11 @@ class ConvertFunctionInst final
   bool withoutActuallyEscaping() const {
     return SILInstruction::Bits.ConvertFunctionInst.WithoutActuallyEscaping;
   }
+            
+  /// Returns `true` if the function conversion is between types with the same
+  /// argument and return types, as well as all other attributes, after substitution,
+  /// such as converting `$<A, B> in (A) -> B for <Int, String>` to `(Int) -> String`.
+  bool onlyConvertsSubstitutions() const;
 };
 
 /// ConvertEscapeToNoEscapeInst - Change the type of a escaping function value

include/swift/SIL/SILInstruction.h

@@ -315,6 +315,65 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     super::visitDestroyValueInst(Destroy);
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  void visitDifferentiableFunctionExtractInst(
+      DifferentiableFunctionExtractInst *dfei) {
+    // If the extractee is the original function, do regular cloning.
+    if (dfei->getExtractee() ==
+        NormalDifferentiableFunctionTypeComponent::Original) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // If the extractee is a derivative function, check whether the *remapped
+    // derivative function type* (BC) is equal to the *derivative remapped
+    // function type* (AD).
+    //
+    // +----------------+        remap       +-------------------------+
+    // | orig.  fn type |  -------(A)------> | remapped orig.  fn type |
+    // +----------------+                    +-------------------------+
+    //         |                                                |
+    //    (B, SILGen)   getAutoDiffDerivativeFunctionType   (D, here)
+    //         V                                                V
+    // +----------------+        remap       +-------------------------+
+    // | deriv. fn type |  -------(C)------> | remapped deriv. fn type |
+    // +----------------+                    +-------------------------+
+    //
+    // (AD) does not always commute with (BC):
+    // - (AD) is the result of remapping, then computing the derivative type.
+    //   This is the default cloning behavior, but may break invariants in the
+    //   initial SIL generated by SILGen.
+    // - (BC) is the result of computing the derivative type (SILGen), then
+    //   remapping. This is the expected type, preserving invariants from
+    //   earlier transforms.
+    //
+    // If (AD) is not equal to (BC), use (BC) as the explicit type.
+    SILType remappedOrigType = getOpType(dfei->getFunctionOperand()->getType());
+    auto remappedOrigFnType = remappedOrigType.castTo<SILFunctionType>();
+    auto derivativeRemappedFnType =
+        remappedOrigFnType
+            ->getAutoDiffDerivativeFunctionType(
+                remappedOrigFnType->getDifferentiabilityParameterIndices(),
+                /*resultIndex*/ 0, dfei->getDerivativeFunctionKind(),
+                getBuilder().getModule().Types,
+                LookUpConformanceInModule(SwiftMod))
+            ->getWithoutDifferentiability();
+    SILType remappedDerivativeFnType = getOpType(dfei->getType());
+    // If remapped derivative type and derivative remapped type are equal, do
+    // regular cloning.
+    if (SILType::getPrimitiveObjectType(derivativeRemappedFnType) ==
+        remappedDerivativeFnType) {
+      super::visitDifferentiableFunctionExtractInst(dfei);
+      return;
+    }
+    // Otherwise, explicitly use the remapped derivative type.
+    recordClonedInstruction(
+        dfei,
+        getBuilder().createDifferentiableFunctionExtract(
+            getOpLocation(dfei->getLoc()), dfei->getExtractee(),
+            getOpValue(dfei->getFunctionOperand()), remappedDerivativeFnType));
+  }
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// One abstract function in the debug info can only have one set of variables
   /// and types. This function determines whether applying the substitutions in
   /// \p SubsMap on the generic signature \p Sig will change the generic type

include/swift/SIL/TypeSubstCloner.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/AutoDiff.h"
 #include "swift/Basic/LLVM.h"
+#include "swift/SIL/SILBuilder.h"
 
 namespace swift {
 
@@ -113,6 +114,13 @@ getOrCreateSubsetParametersThunkForLinearMap(
     AutoDiffDerivativeFunctionKind kind, SILAutoDiffIndices desiredIndices,
     SILAutoDiffIndices actualIndices);
 
+/// Reabstracts the given function-typed value `fn` to the target type `toType`.
+/// Remaps substitutions using `remapSubstitutions`.
+SILValue reabstractFunction(
+    SILBuilder &builder, SILOptFunctionBuilder &fb, SILLocation loc,
+    SILValue fn, CanSILFunctionType toType,
+    std::function<SubstitutionMap(SubstitutionMap)> remapSubstitutions);
+
 } // end namespace autodiff
 
 } // end namespace swift

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

@@ -101,14 +101,17 @@ using AssociativityCacheType =
 
 struct OverrideSignatureKey {
   GenericSignature baseMethodSig;
-  GenericSignature derivedClassSig;
-  Type superclassTy;
+  GenericSignature derivedMethodSig;
+  Type superclassTy, subclassTy;
 
   OverrideSignatureKey(GenericSignature baseMethodSignature,
-                       GenericSignature derivedClassSignature,
-                       Type superclassType)
-      : baseMethodSig(baseMethodSignature),
-        derivedClassSig(derivedClassSignature), superclassTy(superclassType) {}
+                       GenericSignature derivedMethodSignature,
+                       Type superclassType,
+                       Type subclassType)
+    : baseMethodSig(baseMethodSignature),
+      derivedMethodSig(derivedMethodSignature),
+      superclassTy(superclassType),
+      subclassTy(subclassType) {}
 };
 
 namespace llvm {
@@ -119,28 +122,32 @@ template <> struct DenseMapInfo<OverrideSignatureKey> {
   static bool isEqual(const OverrideSignatureKey lhs,
                       const OverrideSignatureKey rhs) {
     return lhs.baseMethodSig.getPointer() == rhs.baseMethodSig.getPointer() &&
-           lhs.derivedClassSig.getPointer() == rhs.derivedClassSig.getPointer() &&
-           lhs.superclassTy.getPointer() == rhs.superclassTy.getPointer();
+           lhs.derivedMethodSig.getPointer() == rhs.derivedMethodSig.getPointer() &&
+           lhs.superclassTy.getPointer() == rhs.superclassTy.getPointer() &&
+           lhs.subclassTy.getPointer() == rhs.subclassTy.getPointer();
   }
 
   static inline OverrideSignatureKey getEmptyKey() {
     return OverrideSignatureKey(DenseMapInfo<GenericSignature>::getEmptyKey(),
                                 DenseMapInfo<GenericSignature>::getEmptyKey(),
+                                DenseMapInfo<Type>::getEmptyKey(),
                                 DenseMapInfo<Type>::getEmptyKey());
   }
 
   static inline OverrideSignatureKey getTombstoneKey() {
     return OverrideSignatureKey(
         DenseMapInfo<GenericSignature>::getTombstoneKey(),
         DenseMapInfo<GenericSignature>::getTombstoneKey(),
+        DenseMapInfo<Type>::getTombstoneKey(),
         DenseMapInfo<Type>::getTombstoneKey());
   }
 
   static unsigned getHashValue(const OverrideSignatureKey &Val) {
     return hash_combine(
         DenseMapInfo<GenericSignature>::getHashValue(Val.baseMethodSig),
-        DenseMapInfo<GenericSignature>::getHashValue(Val.derivedClassSig),
-        DenseMapInfo<Type>::getHashValue(Val.superclassTy));
+        DenseMapInfo<GenericSignature>::getHashValue(Val.derivedMethodSig),
+        DenseMapInfo<Type>::getHashValue(Val.superclassTy),
+        DenseMapInfo<Type>::getHashValue(Val.subclassTy));
   }
 };
 } // namespace llvm
@@ -3563,7 +3570,9 @@ CanSILFunctionType SILFunctionType::get(
   assert(coroutineKind == SILCoroutineKind::None || normalResults.empty());
   assert(coroutineKind != SILCoroutineKind::None || yields.empty());
   assert(!ext.isPseudogeneric() || genericSig);
-
+  
+  substitutions = substitutions.getCanonical();
+  
   llvm::FoldingSetNodeID id;
   SILFunctionType::Profile(id, genericSig, ext, coroutineKind, callee, params,
                            yields, normalResults, errorResult,
@@ -4679,8 +4688,9 @@ ASTContext::getOverrideGenericSignature(const ValueDecl *base,
   unsigned derivedDepth = 0;
 
   auto key = OverrideSignatureKey(baseGenericCtx->getGenericSignature(),
-                                  derivedClass->getGenericSignature(),
-                                  derivedClass->getSuperclass());
+                                  derivedGenericCtx->getGenericSignature(),
+                                  derivedClass->getSuperclass(),
+                                  derivedClass->getDeclaredInterfaceType());
 
   if (getImpl().overrideSigCache.find(key) !=
       getImpl().overrideSigCache.end()) {