[AutoDiff] Use GenericSignature instead of GenericSignatureImpl *. (#27831)

Use `GenericSignature` instead of `GenericSignatureImpl *` in `AutoDiffConfig`
and the differentiation transform.

`AutoDiffConfig` changes by @marcrasi.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -269,11 +269,11 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &s,
 struct AutoDiffConfig {
   IndexSubset *parameterIndices;
   IndexSubset *resultIndices;
-  GenericSignatureImpl* derivativeGenericSignature;
+  GenericSignature derivativeGenericSignature;
 
   /*implicit*/ AutoDiffConfig(IndexSubset *parameterIndices,
                               IndexSubset *resultIndices,
-                              GenericSignatureImpl *derivativeGenericSignature)
+                              GenericSignature derivativeGenericSignature)
       : parameterIndices(parameterIndices), resultIndices(resultIndices),
         derivativeGenericSignature(derivativeGenericSignature) {}
 
@@ -443,7 +443,7 @@ namespace llvm {
 
 using swift::AutoDiffConfig;
 using swift::AutoDiffDerivativeFunctionKind;
-using swift::GenericSignatureImpl;
+using swift::GenericSignature;
 using swift::IndexSubset;
 using swift::SILAutoDiffIndices;
 
@@ -453,27 +453,29 @@ template<> struct DenseMapInfo<AutoDiffConfig> {
   static AutoDiffConfig getEmptyKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getEmptyKey();
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
-            static_cast<GenericSignatureImpl *>(ptr)};
+            DenseMapInfo<GenericSignature>::getEmptyKey()};
   }
 
   static AutoDiffConfig getTombstoneKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
-            static_cast<GenericSignatureImpl *>(ptr)};
+            DenseMapInfo<GenericSignature>::getTombstoneKey()};
   }
 
   static unsigned getHashValue(const AutoDiffConfig &Val) {
     unsigned combinedHash = hash_combine(
         ~1U, DenseMapInfo<void *>::getHashValue(Val.parameterIndices),
         DenseMapInfo<void *>::getHashValue(Val.resultIndices),
-        DenseMapInfo<void *>::getHashValue(Val.derivativeGenericSignature));
+        DenseMapInfo<GenericSignature>::getHashValue(
+            Val.derivativeGenericSignature));
     return combinedHash;
   }
 
   static bool isEqual(const AutoDiffConfig &LHS, const AutoDiffConfig &RHS) {
     return LHS.parameterIndices == RHS.parameterIndices &&
         LHS.resultIndices == RHS.resultIndices &&
-        LHS.derivativeGenericSignature == RHS.derivativeGenericSignature;
+        DenseMapInfo<GenericSignature>::isEqual(LHS.derivativeGenericSignature,
+                                                RHS.derivativeGenericSignature);
   }
 };
 

lib/AST/ASTMangler.cpp

@@ -434,7 +434,7 @@ std::string ASTMangler::mangleSILDifferentiabilityWitnessKey(
   auto originalName = key.first;
   auto *parameterIndices = key.second.parameterIndices;
   auto *resultIndices = key.second.resultIndices;
-  auto *derivativeGenericSignature = key.second.derivativeGenericSignature;
+  auto derivativeGenericSignature = key.second.derivativeGenericSignature;
 
   Buffer << "AD__" << originalName << '_';
   Buffer << "P" << parameterIndices->getString();

lib/SILGen/SILGen.cpp

@@ -810,7 +810,7 @@ void SILGenModule::emitDifferentiabilityWitness(
   bool reorderSelf = shouldReorderSelf();
 
   CanGenericSignature derivativeCanGenSig;
-  if (auto *derivativeGenSig = config.derivativeGenericSignature)
+  if (auto derivativeGenSig = config.derivativeGenericSignature)
     derivativeCanGenSig = derivativeGenSig->getCanonicalSignature();
   // TODO(TF-835): Use simpler derivative generic signature logic below when
   // type-checking no longer generates implicit `@differentiable` attributes.

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1823,18 +1823,18 @@ void LinearMapInfo::generateDifferentiationDataStructures(
 
 class DifferentiableActivityCollection {
 public:
-  SmallDenseMap<GenericSignatureImpl *, DifferentiableActivityInfo> activityInfoMap;
+  SmallDenseMap<GenericSignature, DifferentiableActivityInfo> activityInfoMap;
   SILFunction &function;
   DominanceInfo *domInfo;
   PostDominanceInfo *postDomInfo;
 
   DifferentiableActivityInfo &getActivityInfo(
       GenericSignature assocGenSig, AutoDiffDerivativeFunctionKind kind) {
-    auto activityInfoLookup = activityInfoMap.find(assocGenSig.getPointer());
+    auto activityInfoLookup = activityInfoMap.find(assocGenSig);
     if (activityInfoLookup != activityInfoMap.end())
       return activityInfoLookup->getSecond();
     auto insertion = activityInfoMap.insert(
-        {assocGenSig.getPointer(), DifferentiableActivityInfo(*this, assocGenSig)});
+        {assocGenSig, DifferentiableActivityInfo(*this, assocGenSig)});
     return insertion.first->getSecond();
   }