fix bug where AD is not always seeing custom [differentiable] attrs (#21565)

Author: marcrasi

lib/SILOptimizer/Mandatory/TFDifferentiation.cpp

@@ -1001,6 +1001,7 @@ class ADContext {
   /// parameters as the function.
   StructDecl *createPrimalValueStruct(const DifferentiationTask *task);
 
+private:
   /// Finds the `[differentiable]` attribute on the specified original function
   /// corresponding to the specified parameter indices. Returns nullptr if it
   /// does not exist.
@@ -1078,6 +1079,7 @@ class ADContext {
     return differentiationTasks[existing->getSecond()].get();
   }
 
+public:
   /// Register a differentiation task in the global worklist. This will ensure
   /// that a `[differentiable]` attribute will be generated for the specified
   /// indices, and that primal/adjoint synthesis will be run in the
@@ -1113,6 +1115,21 @@ class ADContext {
   lookUpOrRegisterDifferentiationTask(SILFunction *original,
                                       const SILAutoDiffIndices &indices,
                                       DifferentiationInvoker invoker) {
+    // If `original` has no differentiable attributes, it may be the case that
+    // it has not been loaded yet. Load it, check for differentiable attributes,
+    // and register the attributes as tasks so that they can be looked up.
+    if (original->getDifferentiableAttrs().empty() &&
+        original->isExternalDeclaration()) {
+      auto loaded = module.loadFunction(original);
+      assert(loaded && "Cannot load original function");
+      (void)loaded;
+      for (auto *diffAttr : original->getDifferentiableAttrs()) {
+        registerDifferentiationTask(
+            original, diffAttr->getIndices(),
+            DifferentiationInvoker(diffAttr, original));
+      }
+    }
+
     if (auto *existingTask = lookUpMinimalDifferentiationTask(original, indices))
       return existingTask;
     return registerDifferentiationTask(original, indices, invoker);

test/AutoDiff/superset_adjoint_loaded.swift

@@ -0,0 +1,19 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+// Do not add other tests to this file, because differentiating other things
+// can hide this bug.
+
+// There was a bug where the AD pass would not see the custom [differentiable]
+// attribute on `Float.*` wrt both parameters until after the AD pass generated
+// its own adjoint for `Float.*` wrt the second parameter. This test verifies
+// that the AD pass uses the custom adjoint defined on `Float.*`.
+
+func mul3(_ x: Float) -> Float {
+  return 3 * x
+}
+
+let _ = gradient(at: 0, in: mul3)
+
+// CHECK-LABEL: sil{{.*}} @AD__{{.*}}mul3{{.*}}__primal{{.*}}
+// CHECK: function_ref static Float._adjointMultiply(_:_:_:_:)
+// CHECK: } // end sil function 'AD__{{.*}}mul3{{.*}}__primal{{.*}}'