AutoDiff: Workaround for performing generic signature queries on the wrong signature

Author: slavapestov

lib/SIL/IR/SILFunctionType.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/ProtocolConformance.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILType.h"
@@ -360,6 +361,41 @@ getSemanticResults(SILFunctionType *functionType, IndexSubset *parameterIndices,
       IndexSubset::get(C, parameterIndices->getCapacity(), inoutParamIndices);
 }
 
+static CanGenericSignature buildDifferentiableGenericSignature(CanGenericSignature sig,
+                                                               CanType tanType) {
+  if (!sig)
+    return sig;
+
+  llvm::DenseSet<CanType> types;
+
+  auto &ctx = tanType->getASTContext();
+
+  (void) tanType.findIf([&](Type t) -> bool {
+    if (auto *dmt = t->getAs<DependentMemberType>()) {
+      if (dmt->getName() == ctx.Id_TangentVector)
+        types.insert(dmt->getBase()->getCanonicalType());
+    }
+
+    return false;
+  });
+
+  SmallVector<Requirement, 2> reqs;
+  auto *proto = ctx.getProtocol(KnownProtocolKind::Differentiable);
+  assert(proto != nullptr);
+
+  for (auto type : types) {
+    if (!sig->requiresProtocol(type, proto)) {
+      reqs.push_back(Requirement(RequirementKind::Conformance, type,
+                                 proto->getDeclaredInterfaceType()));
+    }
+  }
+
+  return evaluateOrDefault(
+      ctx.evaluator,
+      AbstractGenericSignatureRequest{sig.getPointer(), {}, reqs},
+      GenericSignature()).getCanonicalSignature();
+}
+
 /// Returns the differential type for the given original function type,
 /// parameter indices, and result index.
 static CanSILFunctionType getAutoDiffDifferentialType(
@@ -371,10 +407,11 @@ static CanSILFunctionType getAutoDiffDifferentialType(
   auto getTangentParameterConvention =
       [&](CanType tanType,
           ParameterConvention origParamConv) -> ParameterConvention {
-    tanType =
-        tanType->getCanonicalType(originalFnTy->getSubstGenericSignature());
-    AbstractionPattern pattern(originalFnTy->getSubstGenericSignature(),
-                               tanType);
+    auto sig = buildDifferentiableGenericSignature(
+      originalFnTy->getSubstGenericSignature(), tanType);
+
+    tanType = tanType->getCanonicalType(sig);
+    AbstractionPattern pattern(sig, tanType);
     auto &tl =
         TC.getTypeLowering(pattern, tanType, TypeExpansionContext::minimal());
     // When the tangent type is address only, we must ensure that the tangent
@@ -398,10 +435,11 @@ static CanSILFunctionType getAutoDiffDifferentialType(
   auto getTangentResultConvention =
       [&](CanType tanType,
           ResultConvention origResConv) -> ResultConvention {
-    tanType =
-        tanType->getCanonicalType(originalFnTy->getSubstGenericSignature());
-    AbstractionPattern pattern(originalFnTy->getSubstGenericSignature(),
-                               tanType);
+    auto sig = buildDifferentiableGenericSignature(
+      originalFnTy->getSubstGenericSignature(), tanType);
+
+    tanType = tanType->getCanonicalType(sig);
+    AbstractionPattern pattern(sig, tanType);
     auto &tl =
         TC.getTypeLowering(pattern, tanType, TypeExpansionContext::minimal());
     // When the tangent type is address only, we must ensure that the tangent
@@ -530,10 +568,11 @@ static CanSILFunctionType getAutoDiffPullbackType(
   auto getTangentParameterConventionForOriginalResult =
       [&](CanType tanType,
           ResultConvention origResConv) -> ParameterConvention {
-    tanType =
-        tanType->getCanonicalType(originalFnTy->getSubstGenericSignature());
-    AbstractionPattern pattern(originalFnTy->getSubstGenericSignature(),
-                               tanType);
+    auto sig = buildDifferentiableGenericSignature(
+      originalFnTy->getSubstGenericSignature(), tanType);
+
+    tanType = tanType->getCanonicalType(sig);
+    AbstractionPattern pattern(sig, tanType);
     auto &tl =
         TC.getTypeLowering(pattern, tanType, TypeExpansionContext::minimal());
     ParameterConvention conv;
@@ -560,10 +599,11 @@ static CanSILFunctionType getAutoDiffPullbackType(
   auto getTangentResultConventionForOriginalParameter =
       [&](CanType tanType,
           ParameterConvention origParamConv) -> ResultConvention {
-    tanType =
-        tanType->getCanonicalType(originalFnTy->getSubstGenericSignature());
-    AbstractionPattern pattern(originalFnTy->getSubstGenericSignature(),
-                               tanType);
+    auto sig = buildDifferentiableGenericSignature(
+      originalFnTy->getSubstGenericSignature(), tanType);
+
+    tanType = tanType->getCanonicalType(sig);
+    AbstractionPattern pattern(sig, tanType);
     auto &tl =
         TC.getTypeLowering(pattern, tanType, TypeExpansionContext::minimal());
     ResultConvention conv;

test/AutoDiff/SIL/Parse/sildeclref.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt %s -module-name=sildeclref_parse | %target-sil-opt -module-name=sildeclref_parse | %FileCheck %s
+// RUN: %target-sil-opt %s -module-name=sildeclref_parse -requirement-machine=off | %target-sil-opt -module-name=sildeclref_parse -requirement-machine=off | %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 | %FileCheck %s
+// RUN: %target-swift-frontend -emit-silgen %s -requirement-machine=off | %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 %s
+// RUN: %target-swift-frontend -emit-sil -verify -requirement-machine=off %s
 
 // Test differentiation transform diagnostics.
 

test/AutoDiff/SILOptimizer/semantic_member_accessors_sil.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -Xllvm -sil-print-after=differentiation %s -module-name null -o /dev/null 2>&1 | %FileCheck %s
+// RUN: %target-swift-frontend -emit-sil -Xllvm -sil-print-after=differentiation %s -module-name null -o /dev/null -requirement-machine=off 2>&1 | %FileCheck %s
 
 // Test differentiation of semantic member accessors:
 // - Stored property accessors.

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

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

test/AutoDiff/mangling.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -enable-experimental-forward-mode-differentiation -module-name=mangling -verify %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-sil -enable-experimental-forward-mode-differentiation -module-name=mangling -verify -requirement-machine=off %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)
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation -Xfrontend -requirement-machine=off)
 // REQUIRES: executable_test
 
 import StdlibUnittest

test/AutoDiff/validation-test/inout_parameters.swift

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