[AutoDiff] Serialize and print @derivative and @transpose accessor kind. (#32839)

Serialize and print the optional accessor kind in `@derivative` and `@transpose`
attributes.

Resolves TF-1293.

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1715,13 +1715,6 @@ class OriginallyDefinedInAttr: public DeclAttribute {
   }
 };
 
-/// A declaration name with location.
-struct DeclNameRefWithLoc {
-  DeclNameRef Name;
-  DeclNameLoc Loc;
-  Optional<AccessorKind> AccessorKind;
-};
-
 /// Attribute that marks a function as differentiable.
 ///
 /// Examples:
@@ -1847,6 +1840,18 @@ class DifferentiableAttr final
   }
 };
 
+/// A declaration name with location.
+struct DeclNameRefWithLoc {
+  /// The declaration name.
+  DeclNameRef Name;
+  /// The declaration name location.
+  DeclNameLoc Loc;
+  /// An optional accessor kind.
+  Optional<AccessorKind> AccessorKind;
+
+  void print(ASTPrinter &Printer) const;
+};
+
 /// The `@derivative(of:)` attribute registers a function as a derivative of
 /// another function-like declaration: a 'func', 'init', 'subscript', or 'var'
 /// computed property declaration.

lib/AST/Attr.cpp

@@ -1052,7 +1052,9 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@derivative");
     Printer << "(of: ";
     auto *attr = cast<DerivativeAttr>(this);
-    Printer << attr->getOriginalFunctionName().Name;
+    if (auto *baseType = attr->getBaseTypeRepr())
+      baseType->print(Printer, Options);
+    attr->getOriginalFunctionName().print(Printer);
     auto *derivative = cast<AbstractFunctionDecl>(D);
     auto diffParamsString = getDifferentiationParametersClauseString(
         derivative, attr->getParameterIndices(), attr->getParsedParameters(),
@@ -1067,7 +1069,9 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@transpose");
     Printer << "(of: ";
     auto *attr = cast<TransposeAttr>(this);
-    Printer << attr->getOriginalFunctionName().Name;
+    if (auto *baseType = attr->getBaseTypeRepr())
+      baseType->print(Printer, Options);
+    attr->getOriginalFunctionName().print(Printer);
     auto *transpose = cast<AbstractFunctionDecl>(D);
     auto transParamsString = getDifferentiationParametersClauseString(
         transpose, attr->getParameterIndices(), attr->getParsedParameters(),
@@ -1719,6 +1723,12 @@ GenericEnvironment *DifferentiableAttr::getDerivativeGenericEnvironment(
   return original->getGenericEnvironment();
 }
 
+void DeclNameRefWithLoc::print(ASTPrinter &Printer) const {
+  Printer << Name;
+  if (AccessorKind)
+    Printer << '.' << getAccessorLabel(*AccessorKind);
+}
+
 void DifferentiableAttr::print(llvm::raw_ostream &OS, const Decl *D,
                                bool omitWrtClause) const {
   StreamPrinter P(OS);

lib/AST/AutoDiff.cpp

@@ -422,6 +422,14 @@ void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
   }
 }
 
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const DeclNameRefWithLoc &name) {
+  os << name.Name;
+  if (auto accessorKind = name.AccessorKind)
+    os << '.' << getAccessorLabel(*accessorKind);
+  return os;
+}
+
 bool swift::operator==(const TangentPropertyInfo::Error &lhs,
                        const TangentPropertyInfo::Error &rhs) {
   if (lhs.kind != rhs.kind)

lib/Serialization/Deserialization.cpp

@@ -4371,16 +4371,26 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
       case decls_block::Derivative_DECL_ATTR: {
         bool isImplicit;
         uint64_t origNameId;
+        bool hasAccessorKind;
+        uint64_t rawAccessorKind;
         DeclID origDeclId;
         uint64_t rawDerivativeKind;
         ArrayRef<uint64_t> parameters;
 
         serialization::decls_block::DerivativeDeclAttrLayout::readRecord(
-            scratch, isImplicit, origNameId, origDeclId, rawDerivativeKind,
-            parameters);
+            scratch, isImplicit, origNameId, hasAccessorKind, rawAccessorKind,
+            origDeclId, rawDerivativeKind, parameters);
+
+        Optional<AccessorKind> accessorKind = None;
+        if (hasAccessorKind) {
+          auto maybeAccessorKind = getActualAccessorKind(rawAccessorKind);
+          if (!maybeAccessorKind)
+            MF.fatal();
+          accessorKind = *maybeAccessorKind;
+        }
 
         DeclNameRefWithLoc origName{DeclNameRef(MF.getDeclBaseName(origNameId)),
-                                    DeclNameLoc(), None};
+                                    DeclNameLoc(), accessorKind};
         auto derivativeKind =
             getActualAutoDiffDerivativeFunctionKind(rawDerivativeKind);
         if (!derivativeKind)

lib/Serialization/ModuleFormat.h

@@ -55,7 +55,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 563; // unchecked_value_cast
+const uint16_t SWIFTMODULE_VERSION_MINOR = 564; // `@derivative` attribute accessor kind
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -1848,6 +1848,8 @@ namespace decls_block {
     Derivative_DECL_ATTR,
     BCFixed<1>, // Implicit flag.
     IdentifierIDField, // Original name.
+    BCFixed<1>, // Has original accessor kind?
+    AccessorKindField, // Original accessor kind.
     DeclIDField, // Original function declaration.
     AutoDiffDerivativeFunctionKindField, // Derivative function kind.
     BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.

lib/Serialization/Serialization.cpp

@@ -2431,19 +2431,25 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       assert(attr->getOriginalFunction(ctx) &&
              "`@derivative` attribute should have original declaration set "
              "during construction or parsing");
-      auto origName = attr->getOriginalFunctionName().Name.getBaseName();
+      auto origDeclNameRef = attr->getOriginalFunctionName();
+      auto origName = origDeclNameRef.Name.getBaseName();
       IdentifierID origNameId = S.addDeclBaseNameRef(origName);
       DeclID origDeclID = S.addDeclRef(attr->getOriginalFunction(ctx));
       auto derivativeKind =
           getRawStableAutoDiffDerivativeFunctionKind(attr->getDerivativeKind());
+      uint8_t rawAccessorKind = 0;
+      auto origAccessorKind = origDeclNameRef.AccessorKind;
+      if (origAccessorKind)
+        rawAccessorKind = uint8_t(getStableAccessorKind(*origAccessorKind));
       auto *parameterIndices = attr->getParameterIndices();
       assert(parameterIndices && "Parameter indices must be resolved");
       SmallVector<bool, 4> paramIndicesVector;
       for (unsigned i : range(parameterIndices->getCapacity()))
         paramIndicesVector.push_back(parameterIndices->contains(i));
       DerivativeDeclAttrLayout::emitRecord(
           S.Out, S.ScratchRecord, abbrCode, attr->isImplicit(), origNameId,
-          origDeclID, derivativeKind, paramIndicesVector);
+          origAccessorKind.hasValue(), rawAccessorKind, origDeclID,
+          derivativeKind, paramIndicesVector);
       return;
     }
 

test/AutoDiff/Serialization/derivative_attr.swift

@@ -56,8 +56,11 @@ extension S {
     (self, { $0 })
   }
 
+  // Note: qualified name base types are not yet serialized and are not printed
+  // when round-tripping.
+
   // CHECK: @derivative(of: instanceMethod, wrt: (self, x))
-  @derivative(of: instanceMethod, wrt: (self, x))
+  @derivative(of: S.instanceMethod, wrt: (self, x))
   func derivativeInstanceMethodWrtAll(_ x: S) -> (value: S, differential: (S, S) -> S) {
     (self, { (dself, dx) in self })
   }
@@ -81,19 +84,39 @@ extension S {
 
 extension S {
   var computedProperty: S {
-    self
+    get { self }
+    set {}
   }
 
   // CHECK: @derivative(of: computedProperty, wrt: self)
   @derivative(of: computedProperty, wrt: self)
   func derivativeProperty() -> (value: S, differential: (S) -> S) {
     (self, { $0 })
   }
+
+  // CHECK: @derivative(of: computedProperty.get, wrt: self)
+  @derivative(of: computedProperty.get, wrt: self)
+  func derivativePropertyGetter() -> (value: S, pullback: (S) -> S) {
+    fatalError()
+  }
+
+  // CHECK: @derivative(of: computedProperty.set, wrt: (self, newValue))
+  @derivative(of: computedProperty.set, wrt: (self, newValue))
+  mutating func derivativePropertySetter(_ newValue: S) -> (
+    value: (), pullback: (inout S) -> S
+  ) {
+    fatalError()
+  }
 }
 
 // Test subscripts.
 
 extension S {
+  subscript() -> S {
+    get { self }
+    set {}
+  }
+
   subscript<T: Differentiable>(x: T) -> S {
     self
   }
@@ -103,4 +126,18 @@ extension S {
   func derivativeSubscript<T: Differentiable>(x: T) -> (value: S, differential: (S) -> S) {
     (self, { $0 })
   }
+
+  // CHECK: @derivative(of: subscript.get, wrt: self)
+  @derivative(of: subscript.get, wrt: self)
+  func derivativeSubscriptGetter() -> (value: S, pullback: (S) -> S) {
+    fatalError()
+  }
+
+  // CHECK: @derivative(of: subscript.set, wrt: (self, newValue))
+  @derivative(of: subscript.set, wrt: (self, newValue))
+  mutating func derivativeSubscriptSetter(_ newValue: S) -> (
+    value: (), pullback: (inout S) -> S
+  ) {
+    fatalError()
+  }
 }

test/AutoDiff/Serialization/transpose_attr.swift

@@ -50,8 +50,11 @@ extension S {
     self + t
   }
 
+  // Note: qualified name base types are not yet serialized and are not printed
+  // when round-tripping.
+
   // CHECK: @transpose(of: instanceMethod, wrt: self)
-  @transpose(of: instanceMethod, wrt: self)
+  @transpose(of: S.instanceMethod, wrt: self)
   static func transposeInstanceMethodWrtSelf(_ other: S, t: S) -> S {
     other + t
   }