[AutoDiff] Lookup for custom derivatives in non-primary source files (#58965)

* Lookup for custom derivatives in non-primary source files after typecheck is finished for the primary source.

This registers all custom derivatives before autodiff transformations and makes them available to them.

Fully resolves #55170

Author: asl

include/swift/AST/Decl.h

@@ -6277,11 +6277,9 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   DerivativeFunctionConfigurationList *DerivativeFunctionConfigs = nullptr;
 
 public:
-  /// Get all derivative function configurations. If `lookInNonPrimarySources`
-  /// is true then lookup is done in non-primary sources as well. Note that
-  /// such lookup might end in cycles if done during sema stages.
+  /// Get all derivative function configurations.
   ArrayRef<AutoDiffConfig>
-  getDerivativeFunctionConfigurations(bool lookInNonPrimarySources = true);
+  getDerivativeFunctionConfigurations();
 
   /// Add the given derivative function configuration.
   void addDerivativeFunctionConfiguration(const AutoDiffConfig &config);

include/swift/Frontend/Frontend.h

@@ -669,6 +669,7 @@ class CompilerInstance {
 
   /// If \p fn returns true, exits early and returns true.
   bool forEachFileToTypeCheck(llvm::function_ref<bool(SourceFile &)> fn);
+  bool forEachSourceFile(llvm::function_ref<bool(SourceFile &)> fn);
 
   /// Whether the cancellation of the current operation has been requested.
   bool isCancellationRequested() const;

include/swift/Subsystems.h

@@ -157,6 +157,10 @@ namespace swift {
   /// emitted.
   void performWholeModuleTypeChecking(SourceFile &SF);
 
+  /// Load derivative configurations from @derivative attributes (including
+  /// those defined in non-primary sources).
+  void loadDerivativeConfigurations(SourceFile &SF);
+
   /// Resolve the given \c TypeRepr to an interface type.
   ///
   /// This is used when dealing with partial source files (e.g. SIL parsing,

lib/AST/Decl.cpp

@@ -8321,7 +8321,7 @@ void AbstractFunctionDecl::prepareDerivativeFunctionConfigurations() {
 }
 
 ArrayRef<AutoDiffConfig>
-AbstractFunctionDecl::getDerivativeFunctionConfigurations(bool lookInNonPrimarySources) {
+AbstractFunctionDecl::getDerivativeFunctionConfigurations() {
   prepareDerivativeFunctionConfigurations();
 
   // Resolve derivative function configurations from `@differentiable`
@@ -8345,36 +8345,6 @@ AbstractFunctionDecl::getDerivativeFunctionConfigurations(bool lookInNonPrimaryS
                                              *DerivativeFunctionConfigs);
   }
 
-  class DerivativeFinder : public ASTWalker {
-    const AbstractFunctionDecl *AFD;
-  public:
-    DerivativeFinder(const AbstractFunctionDecl *afd) : AFD(afd) {}
-
-    bool walkToDeclPre(Decl *D) override {
-      if (auto *afd = dyn_cast<AbstractFunctionDecl>(D)) {
-        for (auto *derAttr : afd->getAttrs().getAttributes<DerivativeAttr>()) {
-          // Resolve derivative function configurations from `@derivative`
-          // attributes by type-checking them.
-          if (AFD->getName().matchesRef(
-                derAttr->getOriginalFunctionName().Name.getFullName())) {
-            (void)derAttr->getOriginalFunction(afd->getASTContext());
-            return false;
-          }
-        }
-      }
-
-      return true;
-    }
-  };
-
-  // Load derivative configurations from @derivative attributes defined in
-  // non-primary sources. Note that it might trigger lookup cycles if called
-  // from inside Sema stages.
-  if (lookInNonPrimarySources) {
-    DerivativeFinder finder(this);
-    getParent()->walkContext(finder);
-  }
-
   return DerivativeFunctionConfigs->getArrayRef();
 }
 

lib/Frontend/Frontend.cpp

@@ -1185,11 +1185,31 @@ bool CompilerInstance::forEachFileToTypeCheck(
   return false;
 }
 
+bool CompilerInstance::forEachSourceFile(
+    llvm::function_ref<bool(SourceFile &)> fn) {
+  for (auto fileName : getMainModule()->getFiles()) {
+    auto *SF = dyn_cast<SourceFile>(fileName);
+    if (!SF) {
+      continue;
+    }
+    if (fn(*SF))
+      return true;
+    ;
+  }
+
+  return false;
+}
+
 void CompilerInstance::finishTypeChecking() {
   forEachFileToTypeCheck([](SourceFile &SF) {
     performWholeModuleTypeChecking(SF);
     return false;
   });
+
+  forEachSourceFile([](SourceFile &SF) {
+    loadDerivativeConfigurations(SF);
+    return false;
+  });
 }
 
 SourceFile::ParsingOptions

lib/Sema/TypeCheckProtocol.cpp

@@ -379,8 +379,7 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
     bool foundExactConfig = false;
     Optional<AutoDiffConfig> supersetConfig = None;
     for (auto witnessConfig :
-           witnessAFD->getDerivativeFunctionConfigurations(
-             /*lookInNonPrimarySources*/ false)) {
+           witnessAFD->getDerivativeFunctionConfigurations()) {
       // All the witness's derivative generic requirements must be satisfied
       // by the requirement's derivative generic requirements OR by the
       // conditional conformance requirements.

lib/Sema/TypeChecker.cpp

@@ -379,6 +379,44 @@ void swift::performWholeModuleTypeChecking(SourceFile &SF) {
   }
 }
 
+void swift::loadDerivativeConfigurations(SourceFile &SF) {
+  if (!isDifferentiableProgrammingEnabled(SF))
+    return;
+
+  auto &Ctx = SF.getASTContext();
+  FrontendStatsTracer tracer(Ctx.Stats,
+                             "load-derivative-configurations");
+
+  class DerivativeFinder : public ASTWalker {
+  public:
+    DerivativeFinder() {}
+
+    bool walkToDeclPre(Decl *D) override {
+      if (auto *afd = dyn_cast<AbstractFunctionDecl>(D)) {
+        for (auto *derAttr : afd->getAttrs().getAttributes<DerivativeAttr>()) {
+          // Resolve derivative function configurations from `@derivative`
+          // attributes by type-checking them.
+          (void)derAttr->getOriginalFunction(D->getASTContext());
+        }
+      }
+
+      return true;
+    }
+  };
+
+  switch (SF.Kind) {
+  case SourceFileKind::Library:
+  case SourceFileKind::Main: {
+    DerivativeFinder finder;
+    SF.walkContext(finder);
+    return;
+  }
+  case SourceFileKind::SIL:
+  case SourceFileKind::Interface:
+    return;
+  }
+}
+
 bool swift::isAdditiveArithmeticConformanceDerivationEnabled(SourceFile &SF) {
   auto &ctx = SF.getASTContext();
   // Return true if `AdditiveArithmetic` derived conformances are explicitly

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/Inputs/derivatives.swift

@@ -0,0 +1,35 @@
+import _Differentiation
+
+@inlinable
+@derivative(of: min)
+func minVJP<T: Comparable & Differentiable>(
+    _ x: T,
+    _ y: T
+) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+    func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+        if x <= y {
+            return (v, .zero)
+        }
+        else {
+            return (.zero, v)
+        }
+    }
+    return (value: min(x, y), pullback: pullback)
+}
+
+@inlinable
+@derivative(of: max)
+func maxVJP<T: Comparable & Differentiable>(
+    _ x: T,
+    _ y: T
+) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+    func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+        if x < y {
+            return (.zero, v)
+        }
+        else {
+            return (v, .zero)
+        }
+    }
+    return (value: max(x, y), pullback: pullback)
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/main.swift

@@ -0,0 +1,9 @@
+// RUN: %target-swift-frontend -emit-sil -verify -primary-file %s %S/Inputs/derivatives.swift -module-name main -o /dev/null
+
+import _Differentiation
+
+@differentiable(reverse)
+func clamp(_ value: Double, _ lowerBound: Double, _ upperBound: Double) -> Double {
+    // No error expected
+    return max(min(value, upperBound), lowerBound)
+}