[AutoDiff] Requestify @differentiable attribute parameter indices. (#28017)

dan-zheng · web-flow · commit 8a172051ad74 · 2019-11-01T15:06:59.000-07:00
Requestify `@differentiable` attribute parameter indices resolution:
`DifferentiableAttributeParameterIndicesRequest`.

This is necessary for type-checking `@differentiable` attributes in
non-primary files. The previous workaround
(`TypeChecker::checkDeclDifferentiableAttributes`) no longer works
because `TypeChecker::validateDecl` has been replaced with
`InterfaceTypeRequest::evaluate`.

Currently, all `@differentiable` attribute type-checking
(`AttributeChecker::visitDifferentiableAttr`) has been moved into
`DifferentiableAttributeParameterIndicesRequest::evaluate`.

In the future, consider splitting
`DifferentiableAttributeParameterIndicesRequest` into multiple requests
for the following functionality:
- `DifferentiableAttr::getJVPFunction`
- `DifferentiableAttr::getVJPFunction`
- `DifferentiableAttr::getDerivativeGenericSignature`
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -1540,6 +1540,7 @@ class DifferentiableAttr final
       private llvm::TrailingObjects<DifferentiableAttr,
                                     ParsedAutoDiffParameter> {
   friend TrailingObjects;
+  friend class DifferentiableAttributeParameterIndicesRequest;
 
   /// The declaration on which the `@differentiable` attribute is declared.
   Decl *OriginalDeclaration = nullptr;
@@ -1558,7 +1559,8 @@ class DifferentiableAttr final
   /// specified.
   FuncDecl *VJPFunction = nullptr;
   /// The differentiation parameters' indices, resolved by the type checker.
-  IndexSubset *ParameterIndices = nullptr;
+  /// The bit stores whether the parameter indices have been computed.
+  llvm::PointerIntPair<IndexSubset *, 1, bool> ParameterIndicesAndBit;
   /// The trailing where clause (optional).
   TrailingWhereClause *WhereClause = nullptr;
   /// The generic signature for autodiff derivative functions. Resolved by the
@@ -1575,9 +1577,9 @@ class DifferentiableAttr final
                               Optional<DeclNameWithLoc> vjp,
                               TrailingWhereClause *clause);
 
-  explicit DifferentiableAttr(Decl *original, bool implicit,
-                              SourceLoc atLoc, SourceRange baseRange,
-                              bool linear, IndexSubset *indices,
+  explicit DifferentiableAttr(Decl *original, bool implicit, SourceLoc atLoc,
+                              SourceRange baseRange, bool linear,
+                              IndexSubset *indices,
                               Optional<DeclNameWithLoc> jvp,
                               Optional<DeclNameWithLoc> vjp,
                               GenericSignature derivativeGenericSignature);
@@ -1611,12 +1613,9 @@ class DifferentiableAttr final
   /// registered VJP.
   Optional<DeclNameWithLoc> getVJP() const { return VJP; }
 
-  IndexSubset *getParameterIndices() const {
-    return ParameterIndices;
-  }
-  void setParameterIndices(IndexSubset *pi) {
-    ParameterIndices = pi;
-  }
+  bool hasComputedParameterIndices() const;
+  IndexSubset *getParameterIndices() const;
+  void setParameterIndices(IndexSubset *paramIndices);
 
   /// The parsed differentiation parameters, i.e. the list of parameters
   /// specified in 'wrt:'.
@@ -1647,8 +1646,7 @@ class DifferentiableAttr final
   void setVJPFunction(FuncDecl *decl);
 
   bool parametersMatch(const DifferentiableAttr &other) const {
-    assert(ParameterIndices && other.ParameterIndices);
-    return ParameterIndices == other.ParameterIndices;
+    return getParameterIndices() == other.getParameterIndices();
   }
 
   /// Get the derivative generic environment for the given `@differentiable`
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -1634,6 +1634,29 @@ class SynthesizeDefaultInitRequest
   bool isCached() const { return true; }
 };
 
+// SWIFT_ENABLE_TENSORFLOW
+class DifferentiableAttributeParameterIndicesRequest :
+    public SimpleRequest<DifferentiableAttributeParameterIndicesRequest,
+                         IndexSubset *(DifferentiableAttr *, Decl *),
+                         CacheKind::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<IndexSubset *>
+  evaluate(Evaluator &evaluator, DifferentiableAttr *attr, Decl *decl) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  Optional<IndexSubset *> getCachedResult() const;
+  void cacheResult(IndexSubset *value) const;
+};
+// SWIFT_ENABLE_TENSORFLOW END
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -31,6 +31,11 @@ SWIFT_REQUEST(TypeChecker, ClassAncestryFlagsRequest,
               AncestryFlags(ClassDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DefaultDefinitionTypeRequest,
               Type(AssociatedTypeDecl *), Cached, NoLocationInfo)
+// SWIFT_ENABLE_TENSORFLOW
+SWIFT_REQUEST(TypeChecker, DifferentiableAttributeParameterIndicesRequest,
+              IndexSubset *(DifferentiableAttr *, Decl *),
+              SeparatelyCached, NoLocationInfo)
+// SWIFT_ENABLE_TENSORFLOW END
 SWIFT_REQUEST(TypeChecker, DefaultTypeRequest,
               Type(KnownProtocolKind, const DeclContext *), SeparatelyCached,
               NoLocationInfo)
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -24,6 +24,8 @@
 #include "swift/AST/GenericSignatureBuilder.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/TypeRepr.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 // SWIFT_ENABLE_TENSORFLOW
 #include "swift/AST/ParameterList.h"
@@ -1461,9 +1463,9 @@ DifferentiableAttr::DifferentiableAttr(Decl *original, bool implicit,
                                        Optional<DeclNameWithLoc> vjp,
                                        GenericSignature derivativeGenSig)
     : DeclAttribute(DAK_Differentiable, atLoc, baseRange, implicit),
-      Linear(linear), JVP(std::move(jvp)), VJP(std::move(vjp)),
-      ParameterIndices(indices) {
+      Linear(linear), JVP(std::move(jvp)), VJP(std::move(vjp)) {
   setOriginalDeclaration(original);
+  setParameterIndices(indices);
   setDerivativeGenericSignature(derivativeGenSig);
 }
 
@@ -1503,6 +1505,31 @@ void DifferentiableAttr::setOriginalDeclaration(Decl *decl) {
   OriginalDeclaration = decl;
 }
 
+bool DifferentiableAttr::hasComputedParameterIndices() const {
+  return ParameterIndicesAndBit.getInt();
+}
+
+IndexSubset *DifferentiableAttr::getParameterIndices() const {
+  assert(getOriginalDeclaration() &&
+         "Original declaration must have been resolved");
+  auto &ctx = getOriginalDeclaration()->getASTContext();
+  return evaluateOrDefault(
+      ctx.evaluator,
+      DifferentiableAttributeParameterIndicesRequest{
+          const_cast<DifferentiableAttr *>(this), getOriginalDeclaration()},
+      nullptr);
+}
+
+void DifferentiableAttr::setParameterIndices(IndexSubset *paramIndices) {
+  assert(getOriginalDeclaration() &&
+         "Original declaration must have been resolved");
+  auto &ctx = getOriginalDeclaration()->getASTContext();
+  ctx.evaluator.cacheOutput(
+      DifferentiableAttributeParameterIndicesRequest{
+          const_cast<DifferentiableAttr *>(this), getOriginalDeclaration()},
+      std::move(paramIndices));
+}
+
 void DifferentiableAttr::setJVPFunction(FuncDecl *decl) {
   JVPFunction = decl;
   if (decl && !JVP)
diff --git a/lib/AST/TypeCheckRequests.cpp b/lib/AST/TypeCheckRequests.cpp
@@ -1072,3 +1072,21 @@ void swift::simple_display(llvm::raw_ostream &out,
   out << "precedence group " << desc.ident << " at ";
   desc.nameLoc.print(out, desc.dc->getASTContext().SourceMgr);
 }
+
+//----------------------------------------------------------------------------//
+// DifferentiableAttributeParameterIndicesRequest computation.
+//----------------------------------------------------------------------------//
+
+Optional<IndexSubset *>
+DifferentiableAttributeParameterIndicesRequest::getCachedResult() const {
+  auto *attr = std::get<0>(getStorage());
+  if (attr->hasComputedParameterIndices())
+    return attr->ParameterIndicesAndBit.getPointer();
+  return None;
+}
+
+void DifferentiableAttributeParameterIndicesRequest::cacheResult(
+    IndexSubset *parameterIndices) const {
+  auto *attr = std::get<0>(getStorage());
+  attr->ParameterIndicesAndBit.setPointerAndInt(parameterIndices, true);
+}
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -3390,6 +3390,7 @@ static void setOriginalFunctionInDifferentiableAttributes(
   for (auto *attr : Attributes.getAttributes<DifferentiableAttr>())
     const_cast<DifferentiableAttr *>(attr)->setOriginalDeclaration(D);
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 /// Parse a single syntactic declaration and return a list of decl
 /// ASTs.  This can return multiple results for var decls that bind to multiple
@@ -3836,6 +3837,7 @@ Parser::parseDecl(ParseDeclOptions Flags,
       // SWIFT_ENABLE_TENSORFLOW END
     }
     // SWIFT_ENABLE_TENSORFLOW
+    // Set original declaration in `@differentiable` attributes.
     setOriginalFunctionInDifferentiableAttributes(D->getAttrs(), D);
     // SWIFT_ENABLE_TENSORFLOW END
   }
@@ -5592,6 +5594,7 @@ Parser::parseDeclVarGetSet(Pattern *pattern, ParseDeclOptions Flags,
   accessors.record(*this, PrimaryVar, Invalid);
 
   // SWIFT_ENABLE_TENSORFLOW
+  // Set original declaration in `@differentiable` attributes.
   for (auto *accessor : accessors.Accessors)
     setOriginalFunctionInDifferentiableAttributes(accessor->getAttrs(),
                                                   accessor);
@@ -5852,6 +5855,7 @@ Parser::parseDeclVar(ParseDeclOptions Flags,
       VD->setStatic(StaticLoc.isValid());
       VD->getAttrs() = Attributes;
       // SWIFT_ENABLE_TENSORFLOW
+      // Set original declaration in `@differentiable` attributes.
       setOriginalFunctionInDifferentiableAttributes(Attributes, VD);
       // SWIFT_ENABLE_TENSORFLOW END
       setLocalDiscriminator(VD);
@@ -7109,6 +7113,7 @@ Parser::parseDeclSubscript(SourceLoc StaticLoc,
   accessors.record(*this, Subscript, (Invalid || !Status.isSuccess()));
 
   // SWIFT_ENABLE_TENSORFLOW
+  // Set original declaration in `@differentiable` attributes.
   for (auto *accessor : accessors.Accessors)
     setOriginalFunctionInDifferentiableAttributes(accessor->getAttrs(),
                                                   accessor);
diff --git a/lib/SIL/SILFunctionBuilder.cpp b/lib/SIL/SILFunctionBuilder.cpp
@@ -83,6 +83,14 @@ void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
       !constant.autoDiffDerivativeFunctionIdentifier &&
       !constant.isStoredPropertyInitializer() &&
       !constant.isThunk()) {
+    // NOTE: Validate `@differentiable` attributes on `AccessorDecl`s by calling
+    // `getParameterIndices`. This is significant to prevent duplicate SIL
+    // `[differentiable]` attribute generation: `getParameterIndices` deletes
+    // `@differentiable` attributes whose original declaration is an
+    // `AbstractStorageDecl`.
+    if (isa<AccessorDecl>(decl))
+      for (auto *A : Attrs.getAttributes<DifferentiableAttr>())
+        (void)A->getParameterIndices();
     for (auto *A : Attrs.getAttributes<DifferentiableAttr>()) {
       // Get lowered argument indices.
       auto *paramIndices = A->getParameterIndices();
diff --git a/lib/Sema/DerivedConformanceDifferentiable.cpp b/lib/Sema/DerivedConformanceDifferentiable.cpp
@@ -613,13 +613,12 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
     // call to the getter.
     if (member->getEffectiveAccess() > AccessLevel::Internal &&
         !member->getAttrs().hasAttribute<DifferentiableAttr>()) {
-      (void)member->getAccessor(AccessorKind::Get)->getInterfaceType();
+      auto *getter = member->getSynthesizedAccessor(AccessorKind::Get);
+      (void)getter->getInterfaceType();
       // If member or its getter already has a `@differentiable` attribute,
       // continue.
       if (member->getAttrs().hasAttribute<DifferentiableAttr>() ||
-          member->getAccessor(AccessorKind::Get)
-              ->getAttrs()
-              .hasAttribute<DifferentiableAttr>())
+          getter->getAttrs().hasAttribute<DifferentiableAttr>())
         continue;
       GenericSignature derivativeGenSig = GenericSignature();
       // If the parent declaration context is an extension, the nominal type may
@@ -628,9 +627,8 @@ getOrSynthesizeTangentVectorStruct(DerivedConformance &derived, Identifier id) {
       if (auto *extDecl = dyn_cast<ExtensionDecl>(parentDC->getAsDecl()))
         derivativeGenSig = extDecl->getGenericSignature();
       auto *diffableAttr = DifferentiableAttr::create(
-          member->getAccessor(AccessorKind::Get), /*implicit*/ true,
-          SourceLoc(), SourceLoc(), /*linear*/ false, {}, None, None,
-          derivativeGenSig);
+          getter, /*implicit*/ true, SourceLoc(), SourceLoc(),
+          /*linear*/ false, {}, None, None, derivativeGenSig);
       member->getAttrs().add(diffableAttr);
       // Set getter `@differentiable` attribute parameter indices.
       diffableAttr->setParameterIndices(IndexSubset::get(C, 1, {0}));
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
diff --git a/lib/Sema/TypeChecker.h b/lib/Sema/TypeChecker.h
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
