Merge pull request swiftlang#39505 from rxwei/82549134-autodiff-requirement-machine-fix-1

Author: swift-ci

lib/SIL/IR/SILFunctionType.cpp

@@ -396,6 +396,32 @@ static CanGenericSignature buildDifferentiableGenericSignature(CanGenericSignatu
       GenericSignature()).getCanonicalSignature();
 }
 
+/// Given an original type, computes its tangent type for the purpose of
+/// building a linear map using this type.  When the original type is an
+/// archetype or contains a type parameter, appends a new generic parameter and
+/// a corresponding replacement type to the given containers.
+static CanType getAutoDiffTangentTypeForLinearMap(
+  Type originalType,
+  LookupConformanceFn lookupConformance,
+  SmallVectorImpl<GenericTypeParamType *> &substGenericParams,
+  SmallVectorImpl<Type> &substReplacements,
+  ASTContext &context
+) {
+  auto maybeTanType = originalType->getAutoDiffTangentSpace(lookupConformance);
+  assert(maybeTanType && "Type does not have a tangent space?");
+  auto tanType = maybeTanType->getCanonicalType();
+  // If concrete, the tangent type is concrete.
+  if (!tanType->hasArchetype() && !tanType->hasTypeParameter())
+    return tanType;
+  // Otherwise, the tangent type is a new generic parameter substituted for the
+  // tangent type.
+  auto gpIndex = substGenericParams.size();
+  auto gpType = CanGenericTypeParamType::get(0, gpIndex, context);
+  substGenericParams.push_back(gpType);
+  substReplacements.push_back(tanType);
+  return gpType;
+}
+
 /// Returns the differential type for the given original function type,
 /// parameter indices, and result index.
 static CanSILFunctionType getAutoDiffDifferentialType(
@@ -471,45 +497,32 @@ static CanSILFunctionType getAutoDiffDifferentialType(
   getDifferentiabilityParameters(originalFnTy, parameterIndices, diffParams);
   SmallVector<SILParameterInfo, 8> differentialParams;
   for (auto &param : diffParams) {
-    auto paramTan =
-        param.getInterfaceType()->getAutoDiffTangentSpace(lookupConformance);
-    assert(paramTan && "Parameter type does not have a tangent space?");
-    auto paramTanType = paramTan->getCanonicalType();
-    auto paramConv = getTangentParameterConvention(paramTanType,
-                                                   param.getConvention());
-    if (!paramTanType->hasArchetype() && !paramTanType->hasTypeParameter()) {
-      differentialParams.push_back(
-          {paramTan->getCanonicalType(), paramConv});
-    } else {
-      auto gpIndex = substGenericParams.size();
-      auto gpType = CanGenericTypeParamType::get(0, gpIndex, ctx);
-      substGenericParams.push_back(gpType);
-      substReplacements.push_back(paramTanType);
-      differentialParams.push_back({gpType, paramConv});
-    }
+    auto paramTanType = getAutoDiffTangentTypeForLinearMap(
+        param.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
+    auto paramConv = getTangentParameterConvention(
+        // FIXME(rdar://82549134): Use `resultTanType` to compute it instead.
+        param.getInterfaceType()
+            ->getAutoDiffTangentSpace(lookupConformance)
+            ->getCanonicalType(),
+        param.getConvention());
+    differentialParams.push_back({paramTanType, paramConv});
   }
   SmallVector<SILResultInfo, 1> differentialResults;
   for (auto resultIndex : resultIndices->getIndices()) {
     // Handle formal original result.
     if (resultIndex < originalFnTy->getNumResults()) {
       auto &result = originalResults[resultIndex];
-      auto resultTan =
-          result.getInterfaceType()->getAutoDiffTangentSpace(lookupConformance);
-      assert(resultTan && "Result type does not have a tangent space?");
-      auto resultTanType = resultTan->getCanonicalType();
-      auto resultConv =
-          getTangentResultConvention(resultTanType, result.getConvention());
-      if (!resultTanType->hasArchetype() &&
-          !resultTanType->hasTypeParameter()) {
-        differentialResults.push_back(
-            {resultTan->getCanonicalType(), resultConv});
-      } else {
-        auto gpIndex = substGenericParams.size();
-        auto gpType = CanGenericTypeParamType::get(0, gpIndex, ctx);
-        substGenericParams.push_back(gpType);
-        substReplacements.push_back(resultTanType);
-        differentialResults.push_back({gpType, resultConv});
-      }
+      auto resultTanType = getAutoDiffTangentTypeForLinearMap(
+          result.getInterfaceType(), lookupConformance,
+          substGenericParams, substReplacements, ctx);
+      auto resultConv = getTangentResultConvention(
+          // FIXME(rdar://82549134): Use `resultTanType` to compute it instead.
+          result.getInterfaceType()
+              ->getAutoDiffTangentSpace(lookupConformance)
+              ->getCanonicalType(),
+          result.getConvention());
+      differentialResults.push_back({resultTanType, resultConv});
       continue;
     }
     // Handle original `inout` parameter.
@@ -524,11 +537,11 @@ static CanSILFunctionType getAutoDiffDifferentialType(
     if (parameterIndices->contains(paramIndex))
       continue;
     auto inoutParam = originalFnTy->getParameters()[paramIndex];
-    auto paramTan = inoutParam.getInterfaceType()->getAutoDiffTangentSpace(
-        lookupConformance);
-    assert(paramTan && "Parameter type does not have a tangent space?");
+    auto inoutParamTanType = getAutoDiffTangentTypeForLinearMap(
+        inoutParam.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
     differentialResults.push_back(
-        {paramTan->getCanonicalType(), ResultConvention::Indirect});
+        {inoutParamTanType, ResultConvention::Indirect});
   }
 
   SubstitutionMap substitutions;
@@ -635,23 +648,16 @@ static CanSILFunctionType getAutoDiffPullbackType(
     // Handle formal original result.
     if (resultIndex < originalFnTy->getNumResults()) {
       auto &origRes = originalResults[resultIndex];
-      auto resultTan = origRes.getInterfaceType()->getAutoDiffTangentSpace(
-          lookupConformance);
-      assert(resultTan && "Result type does not have a tangent space?");
-      auto resultTanType = resultTan->getCanonicalType();
-      auto paramTanConvention = getTangentParameterConventionForOriginalResult(
-          resultTanType, origRes.getConvention());
-      if (!resultTanType->hasArchetype() &&
-          !resultTanType->hasTypeParameter()) {
-        auto resultTanType = resultTan->getCanonicalType();
-        pullbackParams.push_back({resultTanType, paramTanConvention});
-      } else {
-        auto gpIndex = substGenericParams.size();
-        auto gpType = CanGenericTypeParamType::get(0, gpIndex, ctx);
-        substGenericParams.push_back(gpType);
-        substReplacements.push_back(resultTanType);
-        pullbackParams.push_back({gpType, paramTanConvention});
-      }
+      auto resultTanType = getAutoDiffTangentTypeForLinearMap(
+          origRes.getInterfaceType(), lookupConformance,
+          substGenericParams, substReplacements, ctx);
+      auto paramConv = getTangentParameterConventionForOriginalResult(
+          // FIXME(rdar://82549134): Use `resultTanType` to compute it instead.
+          origRes.getInterfaceType()
+              ->getAutoDiffTangentSpace(lookupConformance)
+              ->getCanonicalType(),
+          origRes.getConvention());
+      pullbackParams.push_back({resultTanType, paramConv});
       continue;
     }
     // Handle original `inout` parameter.
@@ -661,28 +667,18 @@ static CanSILFunctionType getAutoDiffPullbackType(
     auto paramIndex =
         std::distance(originalFnTy->getParameters().begin(), &*inoutParamIt);
     auto inoutParam = originalFnTy->getParameters()[paramIndex];
-    auto paramTan = inoutParam.getInterfaceType()->getAutoDiffTangentSpace(
-        lookupConformance);
-    assert(paramTan && "Parameter type does not have a tangent space?");
     // The pullback parameter convention depends on whether the original `inout`
     // paramater is a differentiability parameter.
     // - If yes, the pullback parameter convention is `@inout`.
     // - If no, the pullback parameter convention is `@in_guaranteed`.
+    auto inoutParamTanType = getAutoDiffTangentTypeForLinearMap(
+        inoutParam.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
     bool isWrtInoutParameter = parameterIndices->contains(paramIndex);
     auto paramTanConvention = isWrtInoutParameter
-                                  ? inoutParam.getConvention()
-                                  : ParameterConvention::Indirect_In_Guaranteed;
-    auto paramTanType = paramTan->getCanonicalType();
-    if (!paramTanType->hasArchetype() && !paramTanType->hasTypeParameter()) {
-      pullbackParams.push_back(
-          SILParameterInfo(paramTanType, paramTanConvention));
-    } else {
-      auto gpIndex = substGenericParams.size();
-      auto gpType = CanGenericTypeParamType::get(0, gpIndex, ctx);
-      substGenericParams.push_back(gpType);
-      substReplacements.push_back(paramTanType);
-      pullbackParams.push_back({gpType, paramTanConvention});
-    }
+        ? inoutParam.getConvention()
+        : ParameterConvention::Indirect_In_Guaranteed;
+    pullbackParams.push_back({inoutParamTanType, paramTanConvention});
   }
 
   // Collect pullback results.
@@ -694,21 +690,16 @@ static CanSILFunctionType getAutoDiffPullbackType(
     // and always appear as pullback parameters.
     if (param.isIndirectInOut())
       continue;
-    auto paramTan =
-        param.getInterfaceType()->getAutoDiffTangentSpace(lookupConformance);
-    assert(paramTan && "Parameter type does not have a tangent space?");
-    auto paramTanType = paramTan->getCanonicalType();
+    auto paramTanType = getAutoDiffTangentTypeForLinearMap(
+        param.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
     auto resultTanConvention = getTangentResultConventionForOriginalParameter(
-        paramTanType, param.getConvention());
-    if (!paramTanType->hasArchetype() && !paramTanType->hasTypeParameter()) {
-      pullbackResults.push_back({paramTanType, resultTanConvention});
-    } else {
-      auto gpIndex = substGenericParams.size();
-      auto gpType = CanGenericTypeParamType::get(0, gpIndex, ctx);
-      substGenericParams.push_back(gpType);
-      substReplacements.push_back(paramTanType);
-      pullbackResults.push_back({gpType, resultTanConvention});
-    }
+        // FIXME(rdar://82549134): Use `resultTanType` to compute it instead.
+        param.getInterfaceType()
+            ->getAutoDiffTangentSpace(lookupConformance)
+            ->getCanonicalType(),
+        param.getConvention());
+    pullbackResults.push_back({paramTanType, resultTanConvention});
   }
   SubstitutionMap substitutions;
   if (!substGenericParams.empty()) {

lib/SIL/IR/TypeLowering.cpp

@@ -2599,9 +2599,15 @@ CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
   if (auto *derivativeId = c.getDerivativeFunctionIdentifier()) {
     auto originalFnTy =
         makeConstantInterfaceType(c.asAutoDiffOriginalFunction());
+    // Protocol witness derivatives cannot have a derivative generic signature,
+    // but class method derivatives can.
+    GenericSignature derivativeGenSig = nullptr;
+    if (isa<ClassDecl>(c.getAbstractFunctionDecl()->getInnermostTypeContext()))
+      derivativeGenSig = derivativeId->getDerivativeGenericSignature();
     auto *derivativeFnTy = originalFnTy->getAutoDiffDerivativeFunctionType(
         derivativeId->getParameterIndices(), derivativeId->getKind(),
-        LookUpConformanceInModule(&M));
+        LookUpConformanceInModule(&M),
+        derivativeGenSig);
     return cast<AnyFunctionType>(derivativeFnTy->getCanonicalType());
   }
 

test/AutoDiff/SIL/Parse/sildeclref.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt %s -module-name=sildeclref_parse -requirement-machine=off | %target-sil-opt -module-name=sildeclref_parse -requirement-machine=off | %FileCheck %s
+// RUN: %target-sil-opt %s -module-name=sildeclref_parse | %target-sil-opt -module-name=sildeclref_parse | %FileCheck %s
 // Parse AutoDiff derivative SILDeclRefs via `witness_method` and `class_method` instructions.
 
 import Swift

test/AutoDiff/SILGen/vtable.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-silgen %s -requirement-machine=off | %FileCheck %s
+// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
 
 // Test derivative function vtable entries for `@differentiable` class members:
 // - Methods.

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -verify -requirement-machine=off %s
+// RUN: %target-swift-frontend -emit-sil -requirement-machine=off -verify %s
 
 // Test differentiation transform diagnostics.
 

test/AutoDiff/compiler_crashers_fixed/sr12744-unhandled-pullback-indirect-result.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -verify -requirement-machine=off %s
+// RUN: %target-swift-frontend -emit-sil -verify %s
 
 // SR-12744: Pullback generation crash for unhandled indirect result.
 // May be due to inconsistent derivative function type calculation logic in

test/AutoDiff/compiler_crashers_fixed/sr14240-symbol-in-ir-file-not-tbd-file.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-Xfrontend -requirement-machine=off)
+// RUN: %target-run-simple-swift
 
 // REQUIRES: executable_test
 

test/AutoDiff/mangling.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -enable-experimental-forward-mode-differentiation -module-name=mangling -verify -requirement-machine=off %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-sil -enable-experimental-forward-mode-differentiation -module-name=mangling -verify %s | %FileCheck %s
 
 import _Differentiation
 

test/AutoDiff/validation-test/forward_mode_simple.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation -Xfrontend -requirement-machine=off)
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation)
 // REQUIRES: executable_test
 
 import StdlibUnittest

test/AutoDiff/validation-test/inout_parameters.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-Xfrontend -requirement-machine=off)
+// RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
 // Would fail due to unavailability of swift_autoDiffCreateLinearMapContext.