typecheck all derivative attributes in the module

Author: Marc Rasi

include/swift/AST/ASTContext.h

@@ -311,7 +311,7 @@ class ASTContext final {
   // same parameter indices but different derivative generic signatures.
   llvm::DenseMap<
       std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
-      DerivativeAttr *>
+      llvm::SmallPtrSet<DerivativeAttr *, 1>>
       DerivativeAttrs;
   // SWIFT_ENABLE_TENSORFLOW END
 

lib/SILOptimizer/Utils/Differentiation/Common.cpp

@@ -734,10 +734,6 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
   if (existingWitness)
     return existingWitness;
 
-  assert(original->isExternalDeclaration() &&
-         "SILGen should create differentiability witnesses for all function "
-         "definitions with explicit differentiable attributes");
-
   return SILDifferentiabilityWitness::createDeclaration(
       module, SILLinkage::PublicExternal, original,
       minimalConfig->parameterIndices, minimalConfig->resultIndices,

lib/Sema/TypeCheckAttr.cpp

@@ -19,6 +19,7 @@
 #include "TypeCheckType.h"
 #include "TypeChecker.h"
 #include "swift/AST/ASTVisitor.h"
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/DiagnosticsParse.h"
 #include "swift/AST/GenericEnvironment.h"
@@ -3592,7 +3593,24 @@ DifferentiableAttributeParameterIndicesRequest::evaluate(
 }
 
 // SWIFT_ENABLE_TENSORFLOW
-void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
+/// Typechecks the given derivative attribute `attr` on decl `D`.
+///
+/// Effects are:
+/// - Sets the parameter indices on `attr`.
+/// - Diagnoses errors.
+/// - Stores the attribute in the `ASTContext` list of derivative attributes.
+/// - Stores the derivative configuration in the original function's list of
+///   derivative configurations.
+///
+/// \returns true on error, false on success.
+static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
+                                    DerivativeAttr *attr) {
+
+  // Note: Implementation must be idempotent because it can get called multiple
+  // times for the same attribute.
+
+  auto &diags = Ctx.Diags;
+
   FuncDecl *derivative = cast<FuncDecl>(D);
   auto lookupConformance =
       LookUpConformanceInModule(D->getDeclContext()->getParentModule());
@@ -3612,29 +3630,27 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   auto derivativeResultTupleType = derivativeResultType->getAs<TupleType>();
   if (!derivativeResultTupleType ||
       derivativeResultTupleType->getNumElements() != 2) {
-    diagnose(attr->getLocation(), diag::derivative_attr_expected_result_tuple);
-    attr->setInvalid();
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_expected_result_tuple);
+    return true;
   }
   auto valueResultElt = derivativeResultTupleType->getElement(0);
   auto funcResultElt = derivativeResultTupleType->getElement(1);
   // Get derivative kind and derivative function identifier.
   AutoDiffDerivativeFunctionKind kind;
   if (valueResultElt.getName().str() != "value") {
-    diagnose(attr->getLocation(),
-             diag::derivative_attr_invalid_result_tuple_value_label);
-    attr->setInvalid();
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_invalid_result_tuple_value_label);
+    return true;
   }
   if (funcResultElt.getName().str() == "differential") {
     kind = AutoDiffDerivativeFunctionKind::JVP;
   } else if (funcResultElt.getName().str() == "pullback") {
     kind = AutoDiffDerivativeFunctionKind::VJP;
   } else {
-    diagnose(attr->getLocation(),
-             diag::derivative_attr_invalid_result_tuple_func_label);
-    attr->setInvalid();
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_invalid_result_tuple_func_label);
+    return true;
   }
   attr->setDerivativeKind(kind);
   // `value: R` result tuple element must conform to `Differentiable`.
@@ -3645,11 +3661,10 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   auto valueResultConf = TypeChecker::conformsToProtocol(
       valueResultType, diffableProto, derivative->getDeclContext(), None);
   if (!valueResultConf) {
-    diagnose(attr->getLocation(),
-             diag::derivative_attr_result_value_not_differentiable,
-             valueResultElt.getType());
-    attr->setInvalid();
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_value_not_differentiable,
+                   valueResultElt.getType());
+    return true;
   }
 
   // Compute expected original function type and look up original function.
@@ -3693,23 +3708,23 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   };
 
   auto noneValidDiagnostic = [&]() {
-    diagnose(originalName.Loc,
-             diag::autodiff_attr_original_decl_none_valid_found,
-             originalName.Name, originalFnType);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_none_valid_found,
+                   originalName.Name, originalFnType);
   };
   auto ambiguousDiagnostic = [&]() {
-    diagnose(originalName.Loc, diag::attr_ambiguous_reference_to_decl,
-             originalName.Name, attr->getAttrName());
+    diags.diagnose(originalName.Loc, diag::attr_ambiguous_reference_to_decl,
+                   originalName.Name, attr->getAttrName());
   };
   auto notFunctionDiagnostic = [&]() {
-    diagnose(originalName.Loc,
-             diag::autodiff_attr_original_decl_invalid_kind,
-             originalName.Name);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_invalid_kind,
+                   originalName.Name);
   };
   std::function<void()> invalidTypeContextDiagnostic = [&]() {
-    diagnose(originalName.Loc,
-             diag::autodiff_attr_original_decl_not_same_type_context,
-             originalName.Name);
+    diags.diagnose(originalName.Loc,
+                   diag::autodiff_attr_original_decl_not_same_type_context,
+                   originalName.Name);
   };
 
   // Returns true if the derivative function and original function candidate are
@@ -3743,52 +3758,39 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
       ambiguousDiagnostic, notFunctionDiagnostic, lookupOptions,
       hasValidTypeContext, invalidTypeContextDiagnostic);
   if (!originalAFD) {
-    attr->setInvalid();
-    return;
+    return true;
   }
   // Diagnose original stored properties. Stored properties cannot have custom
   // registered derivatives.
   if (auto *accessorDecl = dyn_cast<AccessorDecl>(originalAFD)) {
     auto *asd = accessorDecl->getStorage();
     if (asd->hasStorage()) {
-      diagnose(originalName.Loc,
-               diag::derivative_attr_original_stored_property_unsupported,
-               originalName.Name);
-      diagnose(originalAFD->getLoc(), diag::decl_declared_here,
-               asd->getFullName());
-      attr->setInvalid();
-      return;
+      diags.diagnose(originalName.Loc,
+                     diag::derivative_attr_original_stored_property_unsupported,
+                     originalName.Name);
+      diags.diagnose(originalAFD->getLoc(), diag::decl_declared_here,
+                     asd->getFullName());
+      return true;
     }
   }
   attr->setOriginalFunction(originalAFD);
 
-  // Get checked wrt param indices.
-  auto *checkedWrtParamIndices = attr->getParameterIndices();
-
   // Get the parsed wrt param indices, which have not yet been checked.
   // This is defined for parsed attributes.
   auto parsedWrtParams = attr->getParsedParameters();
 
-  // If checked wrt param indices are not specified, compute them.
+  auto *checkedWrtParamIndices = computeDifferentiationParameters(
+      parsedWrtParams, derivative, derivative->getGenericEnvironment(),
+      attr->getAttrName(), attr->getLocation());
   if (!checkedWrtParamIndices)
-    checkedWrtParamIndices =
-        computeDifferentiationParameters(parsedWrtParams, derivative,
-                                         derivative->getGenericEnvironment(),
-                                         attr->getAttrName(),
-                                         attr->getLocation());
-  if (!checkedWrtParamIndices) {
-    attr->setInvalid();
-    return;
-  }
+    return true;
 
   // Check if differentiation parameter indices are valid.
   if (checkDifferentiationParameters(
           originalAFD, checkedWrtParamIndices, originalFnType,
           derivative->getGenericEnvironment(), derivative->getModuleContext(),
-          parsedWrtParams, attr->getLocation())) {
-    attr->setInvalid();
-    return;
-  }
+          parsedWrtParams, attr->getLocation()))
+    return true;
 
   // Set the checked differentiation parameter indices in the attribute.
   attr->setParameterIndices(checkedWrtParamIndices);
@@ -3846,25 +3848,25 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
   // Check if differential/pullback type matches expected type.
   if (!actualFuncEltType->isEqual(expectedFuncEltType)) {
     // Emit differential/pullback type mismatch error on attribute.
-    diagnose(attr->getLocation(),
-             diag::derivative_attr_result_func_type_mismatch,
-             funcResultElt.getName(), originalAFD->getFullName());
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_func_type_mismatch,
+                   funcResultElt.getName(), originalAFD->getFullName());
     // Emit note with expected differential/pullback type on actual type
     // location.
     auto *tupleReturnTypeRepr =
         cast<TupleTypeRepr>(derivative->getBodyResultTypeLoc().getTypeRepr());
     auto *funcEltTypeRepr = tupleReturnTypeRepr->getElementType(1);
-    diagnose(funcEltTypeRepr->getStartLoc(),
-             diag::derivative_attr_result_func_type_mismatch_note,
-             funcResultElt.getName(), expectedFuncEltType)
+    diags
+        .diagnose(funcEltTypeRepr->getStartLoc(),
+                  diag::derivative_attr_result_func_type_mismatch_note,
+                  funcResultElt.getName(), expectedFuncEltType)
         .highlight(funcEltTypeRepr->getSourceRange());
     // Emit note showing original function location, if possible.
     if (originalAFD->getLoc().isValid())
-      diagnose(originalAFD->getLoc(),
-               diag::derivative_attr_result_func_original_note,
-               originalAFD->getFullName());
-    attr->setInvalid();
-    return;
+      diags.diagnose(originalAFD->getLoc(),
+                     diag::derivative_attr_result_func_original_note,
+                     originalAFD->getFullName());
+    return true;
   }
 
   // Check that derivative visibility is at least as restricted as original
@@ -3873,29 +3875,42 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
           originalAFD->getFormalAccessScope() &&
       !derivative->getFormalAccessScope().isChildOf(
           originalAFD->getFormalAccessScope())) {
-    diagnoseAndRemoveAttr(attr, diag::derivative_attr_visibility_too_broad);
-    diagnose(originalAFD->getLoc(),
-             diag::derivative_attr_visibility_too_broad_note);
-    return;
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_visibility_too_broad);
+    diags.diagnose(originalAFD->getLoc(),
+                   diag::derivative_attr_visibility_too_broad_note);
+    return true;
   }
 
   // Reject duplicate `@derivative` attributes.
-  auto insertion = Ctx.DerivativeAttrs.try_emplace(
-      {originalAFD, checkedWrtParamIndices, kind}, attr);
-  if (!insertion.second) {
-    diagnoseAndRemoveAttr(attr,
-                          diag::derivative_attr_original_already_has_derivative,
-                          originalAFD->getFullName());
-    diagnose(insertion.first->getSecond()->getLocation(),
-             diag::differentiable_attr_duplicate_note);
-    return;
+  auto &derivativeAttrs =
+      Ctx.DerivativeAttrs[{originalAFD, checkedWrtParamIndices, kind}];
+  derivativeAttrs.insert(attr);
+  if (derivativeAttrs.size() > 1) {
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_original_already_has_derivative,
+                   originalAFD->getFullName());
+    for (auto *duplicateAttr : derivativeAttrs) {
+      if (duplicateAttr == attr)
+        continue;
+      diags.diagnose(duplicateAttr->getLocation(),
+                     diag::differentiable_attr_duplicate_note);
+    }
+    return true;
   }
 
   // Register derivative function configuration.
   auto *resultIndices = IndexSubset::get(Ctx, 1, {0});
   originalAFD->addDerivativeFunctionConfiguration(
       {checkedWrtParamIndices, resultIndices,
        derivative->getGenericSignature()});
+
+  return false;
+}
+
+void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
+  if (typeCheckDerivativeAttr(Ctx, D, attr))
+    attr->setInvalid();
 }
 
 void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
@@ -4527,3 +4542,23 @@ DynamicallyReplacedDeclRequest::evaluate(Evaluator &evaluator,
 
   return nullptr;
 }
+
+// SWIFT_ENABLE_TENSORFLOW
+void TypeChecker::typeCheckDerivativeAttrs(SourceFile &sourceFile) {
+  class DerivativeAttrBindingWalker : public ASTWalker {
+    bool walkToDeclPre(Decl *decl) override {
+      auto f = dyn_cast<AbstractFunctionDecl>(decl);
+      if (!f)
+        return true;
+      for (auto *attr : f->getAttrs())
+        if (auto *da = dyn_cast<DerivativeAttr>(attr))
+          typeCheckDerivativeAttr(f->getASTContext(), f, da);
+      return true;
+    }
+  };
+
+  DiagnosticTransaction diagTxn(sourceFile.getASTContext().Diags);
+  DerivativeAttrBindingWalker walker;
+  sourceFile.getParentModule()->walk(walker);
+  diagTxn.abort();
+}

lib/Sema/TypeChecker.cpp

@@ -385,6 +385,16 @@ TypeCheckSourceFileRequest::evaluate(Evaluator &eval,
     // to work.
     ::bindExtensions(*SF);
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // Type check all `@derivative` attributes in the module. Later, attribute
+    // checking re-checks all `@derivative` attributes in the primary file(s).
+    // This initial checking pass must occur before the re-checking, so that
+    // re-checking can diagnose duplicate attributes using information about all
+    // the derivatives in the module. The differentiation pass also relies on
+    // the information collected by this initial checking pass to see
+    // derivatives defined in the whole module.
+    TypeChecker::typeCheckDerivativeAttrs(*SF);
+
     // Type check the top-level elements of the source file.
     for (auto D : llvm::makeArrayRef(SF->Decls).slice(StartElem)) {
       if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
@@ -405,9 +415,8 @@ TypeCheckSourceFileRequest::evaluate(Evaluator &eval,
   }
 
   // Checking that benefits from having the whole module available.
-  if (!Ctx.TypeCheckerOpts.DelayWholeModuleChecking) {
+  if (!Ctx.TypeCheckerOpts.DelayWholeModuleChecking)
     performWholeModuleTypeChecking(*SF);
-  }
 
   return true;
 }

lib/Sema/TypeChecker.h

@@ -1595,6 +1595,10 @@ class TypeChecker final {
   // SWIFT_ENABLE_TENSORFLOW
   void checkFunctionBodyCompilerEvaluable(AbstractFunctionDecl *D);
 
+  /// Typechecks all derivative attributes in the module that can affect
+  /// derivative configurations of functions in `sourceFile`.
+  static void typeCheckDerivativeAttrs(SourceFile &sourceFile);
+
   /// If an expression references 'self.init' or 'super.init' in an
   /// initializer context, returns the implicit 'self' decl of the constructor.
   /// Otherwise, return nil.

lib/TBDGen/TBDGen.cpp

@@ -67,8 +67,6 @@ void TBDGenVisitor::addSymbol(StringRef name, SymbolKind kind) {
   if (StringSymbols && kind == SymbolKind::GlobalSymbol) {
     auto isNewValue = StringSymbols->insert(mangled).second;
     (void)isNewValue;
-    if (!isNewValue)
-      llvm::dbgs() << mangled << "\n";
     assert(isNewValue && "symbol appears twice");
   }
 }

test/AutoDiff/Inputs/derivative_attr_type_checking/main/main.swift

@@ -11,13 +11,15 @@ func sin(_ x: Float) -> Float {
 func jvpSin(x: @nondiff Float) -> (value: Float, differential: (Float) -> (Float)) {
   return (x, { $0 })
 }
+// expected-error @+2 {{a derivative already exists for 'sin'}}
 // expected-note @+1 {{other attribute declared here}}
 @derivative(of: sin, wrt: x) // ok
 func vjpSinExplicitWrt(x: Float) -> (value: Float, pullback: (Float) -> Float) {
   return (x, { $0 })
 }
 
-// expected-error @+1 {{a derivative already exists for 'sin'}}
+// expected-error @+2 {{a derivative already exists for 'sin'}}
+// expected-note @+1 {{other attribute declared here}}
 @derivative(of: sin)
 func vjpDuplicate(x: Float) -> (value: Float, pullback: (Float) -> Float) {
   return (x, { $0 })
@@ -554,6 +556,7 @@ func dDerivativesHaveDifferentAccessLevels2(_ x: Float) -> (value: Float, pullba
 
 // Check that cross-file and cross-module duplicate derivatives are rejected.
 
+// expected-error @+2 {{a derivative already exists for 'functionDefinedInOtherFile_publicDerivativeInOtherFile'}}
 // expected-note @+1 {{other attribute declared here}}
 @derivative(of: functionDefinedInOtherFile_publicDerivativeInOtherFile)
 public func crossFileDuplicateDerivative2(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {