Implement @derivative attribut serialization.

Fix test/Serialization/derivative_attr.swift.

Author: dan-zheng

include/swift/AST/Attr.def

@@ -513,8 +513,8 @@ DECL_ATTR(differentiable, Differentiable,
   91)
 DECL_ATTR(derivative, Derivative,
   OnFunc | LongAttribute | AllowMultipleAttributes |
-  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove |
-  NotSerialized, 92)
+  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove,
+  92)
 SIMPLE_DECL_ATTR(compilerEvaluable, CompilerEvaluable,
   OnAccessor | OnFunc | OnConstructor | OnSubscript |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove |
@@ -542,8 +542,8 @@ DECL_ATTR(quoted, Quoted,
 // TODO(TF-999): Remove deprecated `@differentiating` attribute.
 DECL_ATTR(differentiating, Differentiating,
   OnFunc | LongAttribute | AllowMultipleAttributes |
-  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove |
-  NotSerialized, 98)
+  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove,
+  98)
 // SWIFT_ENABLE_TENSORFLOW END
 
 #undef TYPE_ATTR

lib/AST/Attr.cpp

@@ -948,8 +948,8 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@derivative");
     Printer << "(of: ";
     auto *attr = cast<DerivativeAttr>(this);
-    auto *derivative = cast<AbstractFunctionDecl>(D);
     Printer << attr->getOriginalFunctionName().Name;
+    auto *derivative = cast<AbstractFunctionDecl>(D);
     auto diffParamsString = getDifferentiationParametersClauseString(
         derivative, attr->getParameterIndices(), attr->getParsedParameters());
     if (!diffParamsString.empty())
@@ -963,8 +963,8 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer.printAttrName("@transpose");
     Printer << '(';
     auto *attr = cast<TransposeAttr>(this);
-    auto *transpose = cast<AbstractFunctionDecl>(D);
     Printer << attr->getOriginalFunctionName().Name;
+    auto *transpose = cast<AbstractFunctionDecl>(D);
     auto transParamsString = getTransposedParametersClauseString(
         transpose, attr->getParameterIndices(), attr->getParsedParameters());
     if (!transParamsString.empty())

lib/Serialization/Deserialization.cpp

@@ -2120,6 +2120,22 @@ getActualReadWriteImplKind(unsigned rawKind) {
   return None;
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+/// Translate from the serialization DifferentiabilityKind enumerators, which
+/// are guaranteed to be stable, to the AST ones.
+static Optional<swift::AutoDiffDerivativeFunctionKind>
+getActualAutoDiffDerivativeFunctionKind(uint8_t raw) {
+  switch (serialization::AutoDiffDerivativeFunctionKind(raw)) {
+#define CASE(ID)                                                               \
+  case serialization::AutoDiffDerivativeFunctionKind::ID:                      \
+    return {swift::AutoDiffDerivativeFunctionKind::ID};
+  CASE(JVP)
+  CASE(VJP)
+#undef CASE
+  }
+  return None;
+}
+
 void ModuleFile::configureStorage(AbstractStorageDecl *decl,
                                   uint8_t rawOpaqueReadOwnership,
                                   uint8_t rawReadImplKind,
@@ -4151,6 +4167,38 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
         break;
       }
 
+      // SWIFT_ENABLE_TENSORFLOW
+      case decls_block::Derivative_DECL_ATTR: {
+        bool isImplicit;
+        uint64_t origNameId;
+        DeclID origDeclId;
+        uint64_t rawDerivativeKind;
+        ArrayRef<uint64_t> parameters;
+
+        serialization::decls_block::DerivativeDeclAttrLayout::readRecord(
+            scratch, isImplicit, origNameId, origDeclId, rawDerivativeKind,
+            parameters);
+
+        DeclNameWithLoc origName{MF.getIdentifier(origNameId), DeclNameLoc()};
+        auto *origDecl = cast<AbstractFunctionDecl>(MF.getDecl(origDeclId));
+        auto derivativeKind =
+            getActualAutoDiffDerivativeFunctionKind(rawDerivativeKind);
+        if (!derivativeKind)
+          MF.fatal();
+        llvm::SmallBitVector parametersBitVector(parameters.size());
+        for (unsigned i : indices(parameters))
+          parametersBitVector[i] = parameters[i];
+        auto *indices = IndexSubset::get(ctx, parametersBitVector);
+
+        auto *derivAttr = DerivativeAttr::create(
+            ctx, isImplicit, SourceLoc(), SourceRange(), origName, indices);
+        derivAttr->setOriginalFunction(origDecl);
+        derivAttr->setDerivativeKind(*derivativeKind);
+        Attr = derivAttr;
+        break;
+      }
+      // SWIFT_ENABLE_TENSORFLOW END
+
       case decls_block::DynamicReplacement_DECL_ATTR: {
         bool isImplicit;
         uint64_t numArgs;

lib/Serialization/ModuleFormat.h

@@ -52,7 +52,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 = 528; // derivative function config table
+const uint16_t SWIFTMODULE_VERSION_MINOR = 529; // `@derivative` serialization
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -232,6 +232,16 @@ enum class DifferentiabilityKind : uint8_t {
 };
 using DifferentiabilityKindField = BCFixed<2>;
 
+// SWIFT_ENABLE_TENSORFLOW
+// These IDs must \em not be renumbered or reordered without incrementing the
+// module version.
+enum class AutoDiffDerivativeFunctionKind : uint8_t {
+  JVP = 0,
+  VJP = 1
+};
+using AutoDiffDerivativeFunctionKindField = BCFixed<1>;
+// SWIFT_ENABLE_TENSORFLOW END
+
 enum class ForeignErrorConventionKind : uint8_t {
   ZeroResult,
   NonZeroResult,
@@ -1772,19 +1782,20 @@ namespace decls_block {
     BCFixed<1>, // Implicit flag.
     IdentifierIDField, // Original name.
     DeclIDField, // Original function declaration.
+    AutoDiffDerivativeFunctionKindField, // Derivative function kind.
     BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.
   >;
 
   // TODO(TF-999): Remove deprecated `@differentiating` attribute.
   using DifferentiatingDeclAttrLayout = DerivativeDeclAttrLayout;
-  
+
   // SWIFT_ENABLE_TENSORFLOW
   using TransposeDeclAttrLayout = BCRecordLayout<
     Transpose_DECL_ATTR,
     BCFixed<1>, // Implicit flag.
     IdentifierIDField, // Original name.
     DeclIDField, // Original function declaration.
-    BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.
+    BCArray<BCFixed<1>> // Transposed parameter indices' bitvector.
   >;
 
 #define SIMPLE_DECL_ATTR(X, CLASS, ...)         \

lib/Serialization/Serialization.cpp

@@ -2029,6 +2029,21 @@ static uint8_t getRawStableVarDeclIntroducer(swift::VarDecl::Introducer intr) {
   llvm_unreachable("bad variable decl introducer kind");
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+/// Translate from the AST differentiability kind enum to the Serialization enum
+/// values, which are guaranteed to be stable.
+static uint8_t getRawStableAutoDiffDerivativeFunctionKind(
+    swift::AutoDiffDerivativeFunctionKind kind) {
+  switch (kind) {
+  case swift::AutoDiffDerivativeFunctionKind::JVP:
+    return uint8_t(serialization::AutoDiffDerivativeFunctionKind::JVP);
+  case swift::AutoDiffDerivativeFunctionKind::VJP:
+    return uint8_t(serialization::AutoDiffDerivativeFunctionKind::VJP);
+  }
+  llvm_unreachable("bad derivative function kind");
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 /// Returns true if the declaration of \p decl depends on \p problemContext
 /// based on lexical nesting.
 ///
@@ -2129,9 +2144,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     case DAK_PrivateImport:
       llvm_unreachable("cannot serialize attribute");
     // SWIFT_ENABLE_TENSORFLOW
-    case DAK_Derivative:
     case DAK_Transpose:
-    case DAK_Differentiating:
       llvm_unreachable("cannot serialize attribute");
     // SWIFT_ENABLE_TENSORFLOW END
 
@@ -2345,6 +2358,30 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       return;
     }
 
+    // SWIFT_ENABLE_TENSORFLOW
+    case DAK_Derivative:
+    case DAK_Differentiating: {
+      auto abbrCode = S.DeclTypeAbbrCodes[DerivativeDeclAttrLayout::Code];
+      auto *attr = cast<DerivativeAttr>(DA);
+      assert(attr->getOriginalFunction() &&
+             "`@derivative` attribute should have original declaration set "
+             "during construction or parsing");
+      auto origName = attr->getOriginalFunctionName().Name.getBaseName();
+      IdentifierID origNameId = S.addDeclBaseNameRef(origName);
+      DeclID origDeclID = S.addDeclRef(attr->getOriginalFunction());
+      auto derivativeKind =
+          getRawStableAutoDiffDerivativeFunctionKind(attr->getDerivativeKind());
+      auto paramIndices = attr->getParameterIndices();
+      assert(paramIndices && "Parameter indices must be resolved");
+      SmallVector<bool, 4> indices;
+      for (unsigned i : range(paramIndices->getCapacity()))
+        indices.push_back(paramIndices->contains(i));
+      DerivativeDeclAttrLayout::emitRecord(
+          S.Out, S.ScratchRecord, abbrCode, attr->isImplicit(), origNameId,
+          origDeclID, derivativeKind, indices);
+      return;
+    }
+
     case DAK_Quoted: {
       auto abbrCode = S.DeclTypeAbbrCodes[QuotedDeclAttrLayout::Code];
       auto attr = cast<QuotedAttr>(DA);
@@ -2354,6 +2391,7 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
                                        S.addDeclRef(attr->getQuoteDecl()));
       return;
     }
+    // SWIFT_ENABLE_TENSORFLOW END
     }
   }
 

test/Serialization/derivative_attr.swift

@@ -7,24 +7,24 @@
 
 // BCANALYZER-NOT: UnknownCode
 
-// CHECK: @differentiable(wrt: x, jvp: jvpAddWrtX)
-// CHECK-NEXT: @differentiable(wrt: (x, y), vjp: vjpAdd)
 func add(x: Float, y: Float) -> Float {
   return x + y
 }
+// CHECK: @derivative(of: add, wrt: x)
 @derivative(of: add, wrt: x)
 func jvpAddWrtX(x: Float, y: Float) -> (value: Float, differential: (Float) -> (Float)) {
   return (x + y, { $0 })
 }
+// CHECK: @derivative(of: add, wrt: (x, y))
 @derivative(of: add)
 func vjpAdd(x: Float, y: Float) -> (value: Float, pullback: (Float) -> (Float, Float)) {
   return (x + y, { ($0, $0) })
 }
 
-// CHECK: @differentiable(wrt: x, vjp: vjpGeneric where T : Differentiable)
 func generic<T : Numeric>(x: T) -> T {
   return x
 }
+// CHECK: @derivative(of: generic, wrt: x)
 @derivative(of: generic)
 func vjpGeneric<T>(x: T) -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
   where T : Numeric, T : Differentiable
@@ -33,21 +33,47 @@ func vjpGeneric<T>(x: T) -> (value: T, pullback: (T.TangentVector) -> T.TangentV
 }
 
 protocol InstanceMethod : Differentiable {
-  // CHECK: @differentiable(wrt: (self, x), vjp: vjpFoo)
   func foo(_ x: Self) -> Self
-  // CHECK: @differentiable(wrt: (self, x), jvp: jvpBarWrt where T == T.TangentVector)
   func bar<T : Differentiable>(_ x: T) -> Self
 }
 extension InstanceMethod {
+  // CHECK: @derivative(of: foo, wrt: (self, x))
   @derivative(of: foo)
-  func vjpFoo(x: Self) -> (value: Self, pullback: (Self.TangentVector) -> (Self.TangentVector, Self.TangentVector)) {
+  func vjpFoo(x: Self) -> (value: Self, pullback: (TangentVector) -> (TangentVector, TangentVector)) {
     return (x, { ($0, $0) })
   }
 
+  // CHECK: @derivative(of: bar, wrt: (self, x))
   @derivative(of: bar, wrt: (self, x))
-  func jvpBarWrt<T : Differentiable>(_ x: T) -> (value: Self, differential: (Self.TangentVector, T) -> Self.TangentVector)
+  func jvpBarWrt<T : Differentiable>(_ x: T) -> (value: Self, differential: (TangentVector, T) -> TangentVector)
     where T == T.TangentVector
   {
     return (self, { dself, dx in dself })
   }
+
+  // CHECK: @derivative(of: bar, wrt: (self, x))
+  @derivative(of: bar, wrt: (self, x))
+  func vjpBarWrt<T : Differentiable>(_ x: T) -> (value: Self, pullback: (TangentVector) -> (TangentVector, T))
+    where T == T.TangentVector
+  {
+    return (self, { v in (v, .zero) })
+  }
+}
+
+// Test deprecated `@differentiating` attribute.
+// For simplicity, `@differentiating` is serialized/deserialized as
+// `@derivative` attribute.
+
+func subtract(x: Float, y: Float) -> Float {
+  return x - y
+}
+// CHECK: @derivative(of: subtract, wrt: x)
+@differentiating(subtract, wrt: x)
+func jvpSubtractWrtX(x: Float, y: Float) -> (value: Float, differential: (Float) -> (Float)) {
+  return (x - y, { $0 })
+}
+// CHECK: @derivative(of: subtract, wrt: (x, y))
+@differentiating(subtract)
+func vjpSubtract(x: Float, y: Float) -> (value: Float, pullback: (Float) -> (Float, Float)) {
+  return (x - y, { ($0, -$0) })
 }