Collapse the differentiation-specific thunk type generation code into the general version

DougGregor · DougGregor · commit 9c134263de7d · 2022-03-11T16:33:13.000-08:00
We had two copies of this code that had drifted apart. Bring them back
together so there is just one place where we compute the type of a
reabstraction thunk.
diff --git a/include/swift/SIL/TypeLowering.h b/include/swift/SIL/TypeLowering.h
@@ -1204,6 +1204,26 @@ CanSILFunctionType getNativeSILFunctionType(
     Optional<SubstitutionMap> reqtSubs = None,
     ProtocolConformanceRef witnessMethodConformance = ProtocolConformanceRef());
 
+/// The thunk kinds used in the differentiation transform.
+enum class DifferentiationThunkKind {
+  /// A reabstraction thunk.
+  ///
+  /// Reabstraction thunks transform a function-typed value to another one with
+  /// different parameter/result abstraction patterns. This is identical to the
+  /// thunks generated by SILGen.
+  Reabstraction,
+
+  /// An index subset thunk.
+  ///
+  /// An index subset thunk is used transform JVP/VJPs into a version that is
+  /// "wrt" fewer differentiation parameters.
+  /// - Differentials of thunked JVPs use zero for non-requested differentiation
+  ///    parameters.
+  /// - Pullbacks of thunked VJPs discard results for non-requested
+  ///   differentiation parameters.
+  IndexSubset
+};
+
 /// Build the type of a function transformation thunk.
 CanSILFunctionType buildSILFunctionThunkType(
     SILFunction *fn,
@@ -1214,7 +1234,8 @@ CanSILFunctionType buildSILFunctionThunkType(
     GenericEnvironment *&genericEnv,
     SubstitutionMap &interfaceSubs,
     CanType &dynamicSelfType,
-    bool withoutActuallyEscaping);
+    bool withoutActuallyEscaping,
+    Optional<DifferentiationThunkKind> differentiationThunkKind = None);
 
 } // namespace swift
 
diff --git a/include/swift/SILOptimizer/Differentiation/Thunk.h b/include/swift/SILOptimizer/Differentiation/Thunk.h
@@ -47,41 +47,6 @@ class ADContext;
 // moved to a shared location.
 //===----------------------------------------------------------------------===//
 
-/// The thunk kinds used in the differentiation transform.
-enum class DifferentiationThunkKind {
-  /// A reabstraction thunk.
-  ///
-  /// Reabstraction thunks transform a function-typed value to another one with
-  /// different parameter/result abstraction patterns. This is identical to the
-  /// thunks generated by SILGen.
-  Reabstraction,
-
-  /// An index subset thunk.
-  ///
-  /// An index subset thunk is used transform JVP/VJPs into a version that is
-  /// "wrt" fewer differentiation parameters.
-  /// - Differentials of thunked JVPs use zero for non-requested differentiation
-  ///    parameters.
-  /// - Pullbacks of thunked VJPs discard results for non-requested
-  ///   differentiation parameters.
-  IndexSubset
-};
-
-CanGenericSignature buildThunkSignature(SILFunction *fn, bool inheritGenericSig,
-                                        OpenedArchetypeType *openedExistential,
-                                        GenericEnvironment *&genericEnv,
-                                        SubstitutionMap &contextSubs,
-                                        SubstitutionMap &interfaceSubs,
-                                        ArchetypeType *&newArchetype);
-
-/// Build the type of a function transformation thunk.
-CanSILFunctionType buildThunkType(SILFunction *fn,
-                                  CanSILFunctionType &sourceType,
-                                  CanSILFunctionType &expectedType,
-                                  GenericEnvironment *&genericEnv,
-                                  SubstitutionMap &interfaceSubs,
-                                  bool withoutActuallyEscaping,
-                                  DifferentiationThunkKind thunkKind);
 
 /// Get or create a reabstraction thunk from `fromType` to `toType`, to be
 /// called in `caller`.
diff --git a/lib/SIL/IR/SILFunctionType.cpp b/lib/SIL/IR/SILFunctionType.cpp
@@ -2509,23 +2509,28 @@ CanSILFunctionType swift::buildSILFunctionThunkType(
     GenericEnvironment *&genericEnv,
     SubstitutionMap &interfaceSubs,
     CanType &dynamicSelfType,
-    bool withoutActuallyEscaping) {
+    bool withoutActuallyEscaping,
+    Optional<DifferentiationThunkKind> differentiationThunkKind) {
   // We shouldn't be thunking generic types here, and substituted function types
   // ought to have their substitutions applied before we get here.
   assert(!expectedType->isPolymorphic() &&
          !expectedType->getCombinedSubstitutions());
   assert(!sourceType->isPolymorphic() &&
          !sourceType->getCombinedSubstitutions());
 
-  // Can't build a thunk without context, so we require ownership semantics
-  // on the result type.
-  assert(expectedType->getExtInfo().hasContext());
-
   // This may inherit @noescape from the expectedType. The @noescape attribute
   // is only stripped when using this type to materialize a new decl.
-  auto extInfoBuilder =
-      expectedType->getExtInfo().intoBuilder().withRepresentation(
+  auto extInfoBuilder = expectedType->getExtInfo().intoBuilder();
+  if (!differentiationThunkKind ||
+      *differentiationThunkKind == DifferentiationThunkKind::Reabstraction ||
+      extInfoBuilder.hasContext()) {
+    // Can't build a reabstraction thunk without context, so we require
+    // ownership semantics on the result type.
+    assert(expectedType->getExtInfo().hasContext());
+
+    extInfoBuilder = extInfoBuilder.withRepresentation(
           SILFunctionType::Representation::Thin);
+  }
 
   if (withoutActuallyEscaping)
     extInfoBuilder = extInfoBuilder.withNoEscape(false);
@@ -2631,10 +2636,14 @@ CanSILFunctionType swift::buildSILFunctionThunkType(
   SmallVector<SILParameterInfo, 4> params;
   params.append(expectedType->getParameters().begin(),
                 expectedType->getParameters().end());
-  params.push_back({sourceType,
-                    sourceType->getExtInfo().hasContext()
-                      ? contextConvention
-                      : ParameterConvention::Direct_Unowned});
+
+  if (!differentiationThunkKind ||
+      *differentiationThunkKind == DifferentiationThunkKind::Reabstraction) {
+    params.push_back({sourceType,
+                      sourceType->getExtInfo().hasContext()
+                        ? contextConvention
+                        : ParameterConvention::Direct_Unowned});
+  }
 
   // If this thunk involves DynamicSelfType in any way, add a capture for it
   // in case we need to recover metadata.
diff --git a/lib/SILOptimizer/Differentiation/Thunk.cpp b/lib/SILOptimizer/Differentiation/Thunk.cpp
@@ -36,220 +36,19 @@ namespace autodiff {
 // moved to a shared location.
 //===----------------------------------------------------------------------===//
 
-CanGenericSignature buildThunkSignature(SILFunction *fn, bool inheritGenericSig,
-                                        OpenedArchetypeType *openedExistential,
-                                        GenericEnvironment *&genericEnv,
-                                        SubstitutionMap &contextSubs,
-                                        SubstitutionMap &interfaceSubs,
-                                        ArchetypeType *&newArchetype) {
-  // If there's no opened existential, we just inherit the generic environment
-  // from the parent function.
-  if (openedExistential == nullptr) {
-    auto genericSig = fn->getLoweredFunctionType()->getInvocationGenericSignature();
-    genericEnv = fn->getGenericEnvironment();
-    interfaceSubs = fn->getForwardingSubstitutionMap();
-    contextSubs = interfaceSubs;
-    return genericSig;
-  }
-
-  auto &ctx = fn->getASTContext();
-
-  // Add the existing generic signature.
-  GenericSignature baseGenericSig;
-  int depth = 0;
-  if (inheritGenericSig) {
-    baseGenericSig = fn->getLoweredFunctionType()->getInvocationGenericSignature();
-    if (baseGenericSig)
-      depth = baseGenericSig.getGenericParams().back()->getDepth() + 1;
-  }
-
-  // Add a new generic parameter to replace the opened existential.
-  auto *newGenericParam =
-      GenericTypeParamType::get(/*type sequence*/ false, depth, 0, ctx);
-
-  assert(openedExistential->isRoot());
-  auto constraint = openedExistential->getExistentialType();
-  if (auto existential = constraint->getAs<ExistentialType>())
-    constraint = existential->getConstraintType();
-
-  Requirement newRequirement(RequirementKind::Conformance, newGenericParam,
-                             constraint);
-
-  auto genericSig = buildGenericSignature(ctx, baseGenericSig,
-                                          { newGenericParam },
-                                          { newRequirement });
-  genericEnv = genericSig.getGenericEnvironment();
-
-  newArchetype =
-      genericEnv->mapTypeIntoContext(newGenericParam)->castTo<ArchetypeType>();
-
-  // Calculate substitutions to map the caller's archetypes to the thunk's
-  // archetypes.
-  if (auto calleeGenericSig =
-          fn->getLoweredFunctionType()->getSubstGenericSignature()) {
-    contextSubs = SubstitutionMap::get(
-        calleeGenericSig,
-        [&](SubstitutableType *type) -> Type {
-          return genericEnv->mapTypeIntoContext(type);
-        },
-        MakeAbstractConformanceForGenericType());
-  }
-
-  // Calculate substitutions to map interface types to the caller's archetypes.
-  interfaceSubs = SubstitutionMap::get(
-      genericSig,
-      [&](SubstitutableType *type) -> Type {
-        if (type->isEqual(newGenericParam))
-          return openedExistential;
-        return fn->mapTypeIntoContext(type);
-      },
-      MakeAbstractConformanceForGenericType());
-
-  return genericSig.getCanonicalSignature();
-}
-
 CanSILFunctionType buildThunkType(SILFunction *fn,
                                   CanSILFunctionType &sourceType,
                                   CanSILFunctionType &expectedType,
                                   GenericEnvironment *&genericEnv,
                                   SubstitutionMap &interfaceSubs,
                                   bool withoutActuallyEscaping,
                                   DifferentiationThunkKind thunkKind) {
-  assert(!expectedType->isPolymorphic() &&
-         !expectedType->getCombinedSubstitutions());
-  assert(!sourceType->isPolymorphic() &&
-         !sourceType->getCombinedSubstitutions());
-
-  // Cannot build a reabstraction thunk without context. Ownership semantics
-  // on the result type are required.
-  if (thunkKind == DifferentiationThunkKind::Reabstraction)
-    assert(expectedType->getExtInfo().hasContext());
-
-  // This may inherit @noescape from the expected type. The `@noescape`
-  // attribute is only stripped when using this type to materialize a new decl.
-  // Use `@convention(thin)` if:
-  // - Building a reabstraction thunk type.
-  // - Building an index subset thunk type, where the expected type has context
-  //   (i.e. is `@convention(thick)`).
-  auto extInfoBuilder = expectedType->getExtInfo().intoBuilder();
-  if (thunkKind == DifferentiationThunkKind::Reabstraction ||
-      extInfoBuilder.hasContext()) {
-    extInfoBuilder = extInfoBuilder.withRepresentation(
-        SILFunctionType::Representation::Thin);
-  }
-  if (withoutActuallyEscaping)
-    extInfoBuilder = extInfoBuilder.withNoEscape(false);
-
-  // Does the thunk type involve archetypes other than opened existentials?
-  bool hasArchetypes = false;
-  // Does the thunk type involve an open existential type?
-  CanOpenedArchetypeType openedExistential;
-  auto archetypeVisitor = [&](CanType t) {
-    if (auto archetypeTy = dyn_cast<ArchetypeType>(t)) {
-      if (auto opened = dyn_cast<OpenedArchetypeType>(archetypeTy)) {
-        const auto root = cast<OpenedArchetypeType>(CanType(opened->getRoot()));
-        assert((openedExistential == CanArchetypeType() ||
-                openedExistential == root) &&
-               "one too many open existentials");
-        openedExistential = root;
-      } else {
-        hasArchetypes = true;
-      }
-    }
-  };
-
-  // Use the generic signature from the context if the thunk involves
-  // generic parameters.
-  CanGenericSignature genericSig;
-  SubstitutionMap contextSubs;
-  ArchetypeType *newArchetype = nullptr;
-
-  if (expectedType->hasArchetype() || sourceType->hasArchetype()) {
-    expectedType.visit(archetypeVisitor);
-    sourceType.visit(archetypeVisitor);
-    genericSig =
-        buildThunkSignature(fn, hasArchetypes, openedExistential, genericEnv,
-                            contextSubs, interfaceSubs, newArchetype);
-  }
-
-  auto substTypeHelper = [&](SubstitutableType *type) -> Type {
-    if (CanType(type) == openedExistential)
-      return newArchetype;
-    return Type(type).subst(contextSubs);
-  };
-  auto substConformanceHelper = LookUpConformanceInSubstitutionMap(contextSubs);
-
-  // Utility function to apply contextSubs, and also replace the
-  // opened existential with the new archetype.
-  auto substLoweredTypeIntoThunkContext =
-      [&](CanSILFunctionType t) -> CanSILFunctionType {
-    return SILType::getPrimitiveObjectType(t)
-        .subst(fn->getModule(), substTypeHelper, substConformanceHelper)
-        .castTo<SILFunctionType>();
-  };
-
-  sourceType = substLoweredTypeIntoThunkContext(sourceType);
-  expectedType = substLoweredTypeIntoThunkContext(expectedType);
-
-  // If our parent function was pseudogeneric, this thunk must also be
-  // pseudogeneric, since we have no way to pass generic parameters.
-  if (genericSig)
-    if (fn->getLoweredFunctionType()->isPseudogeneric())
-      extInfoBuilder = extInfoBuilder.withIsPseudogeneric();
-
-  // Add the function type as the parameter.
-  auto contextConvention =
-      SILType::getPrimitiveObjectType(sourceType).isTrivial(*fn)
-          ? ParameterConvention::Direct_Unowned
-          : ParameterConvention::Direct_Guaranteed;
-  SmallVector<SILParameterInfo, 4> params;
-  params.append(expectedType->getParameters().begin(),
-                expectedType->getParameters().end());
-  // Add reabstraction function parameter only if building a reabstraction thunk
-  // type.
-  if (thunkKind == DifferentiationThunkKind::Reabstraction)
-    params.push_back({sourceType, sourceType->getExtInfo().hasContext()
-                                      ? contextConvention
-                                      : ParameterConvention::Direct_Unowned});
-
-  auto mapTypeOutOfContext = [&](CanType type) -> CanType {
-    return type->mapTypeOutOfContext()->getCanonicalType(genericSig);
-  };
-
-  // Map the parameter and expected types out of context to get the interface
-  // type of the thunk.
-  SmallVector<SILParameterInfo, 4> interfaceParams;
-  interfaceParams.reserve(params.size());
-  for (auto &param : params) {
-    auto interfaceParam = param.map(mapTypeOutOfContext);
-    interfaceParams.push_back(interfaceParam);
-  }
-
-  SmallVector<SILYieldInfo, 4> interfaceYields;
-  for (auto &yield : expectedType->getYields()) {
-    auto interfaceYield = yield.map(mapTypeOutOfContext);
-    interfaceYields.push_back(interfaceYield);
-  }
-
-  SmallVector<SILResultInfo, 4> interfaceResults;
-  for (auto &result : expectedType->getResults()) {
-    auto interfaceResult = result.map(mapTypeOutOfContext);
-    interfaceResults.push_back(interfaceResult);
-  }
-
-  Optional<SILResultInfo> interfaceErrorResult;
-  if (expectedType->hasErrorResult()) {
-    auto errorResult = expectedType->getErrorResult();
-    interfaceErrorResult = errorResult.map(mapTypeOutOfContext);
-  }
-
-  // The type of the thunk function.
-  return SILFunctionType::get(
-      genericSig, extInfoBuilder.build(), expectedType->getCoroutineKind(),
-      ParameterConvention::Direct_Unowned, interfaceParams, interfaceYields,
-      interfaceResults, interfaceErrorResult,
-      expectedType->getPatternSubstitutions(), SubstitutionMap(),
-      fn->getASTContext());
+  CanType inputSubstType;
+  CanType outputSubstType;
+  CanType dynamicSelfType;
+  return buildSILFunctionThunkType(
+      fn, sourceType, expectedType, inputSubstType, outputSubstType, genericEnv,
+      interfaceSubs, dynamicSelfType, withoutActuallyEscaping, thunkKind);
 }
 
 /// Forward function arguments, handling ownership convention mismatches.
