Merge branch 'tensorflow' of https://github.com/apple/swift into tensorflow

Author: bartchr808

docs/ABI/Mangling.rst

@@ -516,6 +516,10 @@ Types
   FUNCTION-KIND ::= 'C'                      // C function pointer type
   FUNCTION-KIND ::= 'A'                      // @auto_closure function type (escaping)
   FUNCTION-KIND ::= 'E'                      // function type (noescape)
+  FUNCTION-KIND ::= 'F'                      // @differentiable function type
+  FUNCTION-KIND ::= 'G'                      // @differentiable function type (escaping)
+  FUNCTION-KIND ::= 'H'                      // @differentiable(linear) function type
+  FUNCTION-KIND ::= 'I'                      // @differentiable(linear) function type (escaping)
 
   function-signature ::= params-type params-type throws? // results and parameters
 

include/swift/ABI/MetadataValues.h

@@ -733,6 +733,14 @@ enum class FunctionMetadataConvention: uint8_t {
   CFunctionPointer = 3,
 };
 
+/// Differentiability kind for function type metadata.
+/// Duplicates `DifferentiabilityKind` in AutoDiff.h.
+enum class FunctionMetadataDifferentiabilityKind: uint8_t {
+  NonDifferentiable = 0b00,
+  Normal = 0b01,
+  Linear = 0b11
+};
+
 /// Flags in a function type metadata record.
 template <typename int_type>
 class TargetFunctionTypeFlags {
@@ -747,7 +755,8 @@ class TargetFunctionTypeFlags {
     ParamFlagsMask    = 0x02000000U,
     EscapingMask      = 0x04000000U,
     // SWIFT_ENABLE_TENSORFLOW
-    DifferentiableMask  = 0x08000000U
+    DifferentiableMask  = 0x08000000U,
+    LinearMask          = 0x10000000U
   };
   int_type Data;
 
@@ -785,10 +794,14 @@ class TargetFunctionTypeFlags {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  constexpr TargetFunctionTypeFlags<int_type>
-  withDifferentiable(bool isDifferentiable) const {
-    return TargetFunctionTypeFlags<int_type>((Data & ~DifferentiableMask) |
-                                   (isDifferentiable ? DifferentiableMask : 0));
+  constexpr TargetFunctionTypeFlags<int_type> withDifferentiabilityKind(
+      FunctionMetadataDifferentiabilityKind differentiability) const {
+    return TargetFunctionTypeFlags<int_type>(
+        (Data & ~DifferentiableMask & ~LinearMask) |
+        (differentiability == FunctionMetadataDifferentiabilityKind::Normal
+             ? DifferentiableMask : 0) |
+        (differentiability == FunctionMetadataDifferentiabilityKind::Linear
+             ? LinearMask : 0));
   }
 
   unsigned getNumParameters() const { return Data & NumParametersMask; }
@@ -807,7 +820,15 @@ class TargetFunctionTypeFlags {
 
   // SWIFT_ENABLE_TENSORFLOW
   bool isDifferentiable() const {
-    return bool (Data & DifferentiableMask);
+    return getDifferentiabilityKind() >=
+        FunctionMetadataDifferentiabilityKind::Normal;
+  }
+  FunctionMetadataDifferentiabilityKind getDifferentiabilityKind() const {
+    if (bool(Data & DifferentiableMask))
+      return FunctionMetadataDifferentiabilityKind::Normal;
+    if (bool(Data & LinearMask))
+      return FunctionMetadataDifferentiabilityKind::Linear;
+    return FunctionMetadataDifferentiabilityKind::NonDifferentiable;
   }
 
   bool hasParameterFlags() const { return bool(Data & ParamFlagsMask); }

include/swift/AST/Attr.def

@@ -432,6 +432,9 @@ SIMPLE_DECL_ATTR(compilerEvaluable, CompilerEvaluable,
   /* Not serialized */ 91)
 SIMPLE_DECL_ATTR(noDerivative, NoDerivative,
   OnVar, 92)
+DECL_ATTR(transposing, Transposing,
+  OnFunc | LongAttribute | AllowMultipleAttributes |
+  NotSerialized, 93)
 
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS

include/swift/AST/Attr.h

@@ -1682,6 +1682,79 @@ class DifferentiatingAttr final
     return DA->getKind() == DAK_Differentiating;
   }
 };
+  
+/// Attribute that registers a function as a transpose of another function.
+///
+/// Examples:
+///   @transposing(foo)
+///   @transposing(+, wrt: (lhs, rhs))
+class TransposingAttr final
+      : public DeclAttribute,
+        private llvm::TrailingObjects<DifferentiableAttr,
+                                      ParsedAutoDiffParameter> {
+  /// The base type of the original function.
+  /// This is non-null only when the original function is not top-level (i.e. it
+  /// is an instance/static method).
+  TypeRepr *BaseType;
+  /// The original function name.
+  DeclNameWithLoc Original;
+  /// The original function, resolved by the type checker.
+  FuncDecl *OriginalFunction = nullptr;
+  /// The number of parsed parameters specified in 'wrt:'.
+  unsigned NumParsedParameters = 0;
+  /// The differentiation parameters' indices, resolved by the type checker.
+  AutoDiffIndexSubset *ParameterIndexSubset = nullptr;
+  
+  explicit TransposingAttr(ASTContext &context, bool implicit,
+                           SourceLoc atLoc, SourceRange baseRange,
+                           TypeRepr *baseType, DeclNameWithLoc original,
+                           ArrayRef<ParsedAutoDiffParameter> params);
+  
+  explicit TransposingAttr(ASTContext &context, bool implicit,
+                           SourceLoc atLoc, SourceRange baseRange,
+                           TypeRepr *baseType, DeclNameWithLoc original,
+                           AutoDiffIndexSubset *indices);
+  
+public:
+  static TransposingAttr *create(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 ArrayRef<ParsedAutoDiffParameter> params);
+  
+  static TransposingAttr *create(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 AutoDiffIndexSubset *indices);
+  
+  TypeRepr *getBaseType() const { return BaseType; }
+  DeclNameWithLoc getOriginal() const { return Original; }
+  
+  FuncDecl *getOriginalFunction() const { return OriginalFunction; }
+  void setOriginalFunction(FuncDecl *decl) { OriginalFunction = decl; }
+  
+  /// The parsed transposing parameters, i.e. the list of parameters
+  /// specified in 'wrt:'.
+  ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  MutableArrayRef<ParsedAutoDiffParameter> getParsedParameters() {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  size_t numTrailingObjects(OverloadToken<ParsedAutoDiffParameter>) const {
+    return NumParsedParameters;
+  }
+  
+  AutoDiffIndexSubset *getParameterIndexSubset() const {
+    return ParameterIndexSubset;
+  }
+  void setParameterIndices(AutoDiffIndexSubset *pi) {
+    ParameterIndexSubset = pi;
+  }
+  
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DAK_Transposing;
+  }
+};
 
 /// Attributes that may be applied to declarations.
 class DeclAttributes {

include/swift/AST/AutoDiff.h

@@ -23,7 +23,7 @@
 #include "swift/Basic/Range.h"
 
 namespace swift {
-  
+
 enum class DifferentiabilityKind: uint8_t {
   NonDifferentiable = 0b00,
   Normal = 0b01,
@@ -354,6 +354,13 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
   unsigned getNumIndices() const {
     return (unsigned)std::distance(begin(), end());
   }
+  
+  SmallBitVector getBitVector() const {
+    SmallBitVector indicesBitVec(capacity, false);
+    for (auto index : getIndices())
+      indicesBitVec.set(index);
+    return indicesBitVec;
+  }
 
   bool contains(unsigned index) const {
     unsigned bitWordIndex, offset;

include/swift/AST/DiagnosticsParse.def

@@ -1512,13 +1512,23 @@ ERROR(attr_differentiating_expected_original_name,PointsToFirstBadToken,
 ERROR(attr_differentiating_expected_label_linear_or_wrt,none,
       "expected either 'linear' or 'wrt:'", ())
 
+// transposing
+ERROR(attr_transposing_expected_original_name,PointsToFirstBadToken,
+      "expected an original function name", ())
+ERROR(attr_transposing_expected_label_linear_or_wrt,none,
+      "expected 'wrt:'", ())
+
+// transposing `wrt` parameters clause
+ERROR(transposing_params_clause_expected_parameter,PointsToFirstBadToken,
+      "expected a parameter, which can be a 'unsigned int' parameter number "
+      "or 'self'", ())
+
 // differentiation `wrt` parameters clause
 ERROR(expected_colon_after_label,PointsToFirstBadToken,
       "expected a colon ':' after '%0'", (StringRef))
 ERROR(diff_params_clause_expected_parameter,PointsToFirstBadToken,
       "expected a parameter, which can be a function parameter name, "
-      "parameter index, or 'self'",
-      ())
+      "parameter index, or 'self'", ())
 
 // [differentiable ...] (sil-decl attr)
 ERROR(sil_attr_differentiable_expected_keyword,PointsToFirstBadToken,

include/swift/AST/DiagnosticsSema.def

@@ -2763,8 +2763,10 @@ ERROR(differentiable_attr_overload_not_found,none,
       "%0 does not have expected type %1", (DeclName, Type))
 ERROR(differentiable_attr_no_currying,none,
       "'@differentiable' cannot be defined on functions that return functions", ())
+// TODO(TF-482): Change duplicate `@differentiable` attribute diagnostic to also
+// mention "same generic requirements".
 ERROR(differentiable_attr_duplicate,none,
-      "duplicate '@differentiable' attribute with same parameter indices", ())
+      "duplicate '@differentiable' attribute with same parameters", ())
 NOTE(differentiable_attr_duplicate_note,none,
      "other attribute declared here", ())
 ERROR(differentiable_attr_function_not_same_type_context,none,
@@ -2824,6 +2826,19 @@ ERROR(differentiating_attr_not_in_same_file_as_original,none,
 ERROR(differentiating_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
 
+// transposing
+ERROR(transpose_params_clause_param_not_differentiable,none,
+      "can only transpose with respect to parameters that conform to "
+      "'Differentiable' and where '%0 == %0.TangentVector'", (StringRef))
+ERROR(transposing_attr_overload_not_found,none,
+      "could not find function %0 with expected type %1", (DeclName, Type))
+ERROR(transposing_attr_cant_use_named_wrt_params,none,
+      "cannot use named wrt parameters in '@transposing' attribute, found %0",
+      (Identifier))
+ERROR(transposing_attr_result_value_not_differentiable,none,
+      "'@transposing' attribute requires original function result to "
+      "conform to 'Differentiable'", (Type))
+
 // differentiation `wrt` parameters clause
 ERROR(diff_function_no_parameters,none,
       "%0 has no parameters to differentiate with respect to", (DeclName))
@@ -2841,6 +2856,8 @@ ERROR(diff_params_clause_no_inferred_parameters,PointsToFirstBadToken,
       "no differentiation parameters could be inferred; must differentiate "
       "with respect to at least one parameter conforming to 'Differentiable'",
       ())
+ERROR(diff_params_clause_inout_argument,none,
+      "'inout' parameters (%0) cannot be differentiated with respect to", (Type))
 ERROR(diff_params_clause_cannot_diff_wrt_objects_or_existentials,none,
       "class objects and protocol existentials (%0) cannot be differentiated "
       "with respect to", (Type))
@@ -2865,9 +2882,9 @@ ERROR(compiler_evaluable_ref_non_compiler_evaluable,none,
       "@compilerEvaluable functions may not reference non-@compilerEvaluable functions", ())
 
 // @noDerivative attribute
-ERROR(noderivative_only_on_stored_properties_in_differentiable_structs,none,
-      "'@noDerivative' is only allowed on stored properties in structure types "
-      "that declare a conformance to 'Differentiable'", ())
+ERROR(noderivative_only_on_differentiable_struct_or_class_fields,none,
+      "'@noDerivative' is only allowed on stored properties in structure or "
+      "class types that declare a conformance to 'Differentiable'", ())
 
 //------------------------------------------------------------------------------
 // MARK: Type Check Expressions

include/swift/AST/Types.h

@@ -19,6 +19,7 @@
 
 // SWIFT_ENABLE_TENSORFLOW
 #include "swift/AST/AutoDiff.h"
+#include "swift/AST/Attr.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/GenericParamKey.h"
 #include "swift/AST/Identifier.h"
@@ -274,7 +275,7 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   TypeBase(const TypeBase&) = delete;
   void operator=(const TypeBase&) = delete;
 
-  /// This union contains to the ASTContext for canonical types, and is
+  /// This union contains the ASTContext for canonical types, and is
   /// otherwise lazily populated by ASTContext when the canonical form of a
   /// non-canonical type is requested. The disposition of the union is stored
   /// outside of the union for performance. See Bits.TypeBase.IsCanonical.
@@ -3093,7 +3094,7 @@ class AnyFunctionType : public TypeBase {
   ///
   /// If `makeSelfParamFirst` is true, self's tangent is reordered to appear
   /// first. This should be used during type-checking, e.g. type-checking
-  /// `@differentiable` and `@differentiating` attributes.
+  /// `@differentiable`, `@differentiating`, and `@transposing` attributes.
   ///
   /// \note The original function type (`self`) need not be `@differentiable`.
   /// The resulting function will preserve all `ExtInfo` of the original
@@ -3108,6 +3109,13 @@ class AnyFunctionType : public TypeBase {
   /// Given the type of an autodiff associated function, returns the
   /// corresponding original function type.
   AnyFunctionType *getAutoDiffOriginalFunctionType();
+  
+  /// Given the type of a transposing associated function, returns the
+  /// corresponding original function type.
+  AnyFunctionType *
+  getTransposeOriginalFunctionType(TransposingAttr *attr,
+                                   AutoDiffIndexSubset *wrtParamIndices,
+                                   bool wrtSelf);
 
   AnyFunctionType *getWithoutDifferentiability() const;
 

include/swift/Demangling/DemangleNodes.def

@@ -68,6 +68,12 @@ NODE(DependentProtocolConformanceInherited)
 NODE(DependentProtocolConformanceAssociated)
 CONTEXT_NODE(Destructor)
 CONTEXT_NODE(DidSet)
+// SWIFT_ENABLE_TENSORFLOW
+NODE(DifferentiableFunctionType)
+NODE(EscapingDifferentiableFunctionType)
+NODE(LinearFunctionType)
+NODE(EscapingLinearFunctionType)
+// SWIFT_ENABLE_TENSORFLOW END
 NODE(Directness)
 NODE(DynamicAttribute)
 NODE(DirectMethodReferenceAttribute)

include/swift/Demangling/TypeDecoder.h

@@ -494,6 +494,12 @@ class TypeDecoder {
     case NodeKind::NoEscapeFunctionType:
     case NodeKind::AutoClosureType:
     case NodeKind::EscapingAutoClosureType:
+    // SWIFT_ENABLE_TENSORFLOW
+    case NodeKind::DifferentiableFunctionType:
+    case NodeKind::EscapingDifferentiableFunctionType:
+    case NodeKind::LinearFunctionType:
+    case NodeKind::EscapingLinearFunctionType:
+    // SWIFT_ENABLE_TENSORFLOW END
     case NodeKind::FunctionType: {
       if (Node->getNumChildren() < 2)
         return BuiltType();
@@ -508,6 +514,17 @@ class TypeDecoder {
       } else if (Node->getKind() == NodeKind::ThinFunctionType) {
         flags = flags.withConvention(FunctionMetadataConvention::Thin);
       }
+      // SWIFT_ENABLE_TENSORFLOW
+      else if (Node->getKind() == NodeKind::DifferentiableFunctionType ||
+               Node->getKind() ==
+                   NodeKind::EscapingDifferentiableFunctionType) {
+        flags = flags.withDifferentiabilityKind(
+            FunctionMetadataDifferentiabilityKind::Normal);
+      } else if (Node->getKind() == NodeKind::LinearFunctionType ||
+                 Node->getKind() == NodeKind::EscapingLinearFunctionType) {
+        flags = flags.withDifferentiabilityKind(
+            FunctionMetadataDifferentiabilityKind::Linear);
+      }
 
       bool isThrow =
         Node->getChild(0)->getKind() == NodeKind::ThrowsAnnotation;
@@ -527,7 +544,12 @@ class TypeDecoder {
               .withEscaping(
                           Node->getKind() == NodeKind::FunctionType ||
                           Node->getKind() == NodeKind::EscapingAutoClosureType ||
-                          Node->getKind() == NodeKind::EscapingObjCBlock);
+                          Node->getKind() == NodeKind::EscapingObjCBlock ||
+                          // SWIFT_ENABLE_TENSORFLOW
+                          Node->getKind() ==
+                              NodeKind::EscapingDifferentiableFunctionType ||
+                          Node->getKind() ==
+                              NodeKind::EscapingLinearFunctionType);
 
       auto result = decodeMangledType(Node->getChild(isThrow ? 2 : 1));
       if (!result) return BuiltType();

include/swift/Parse/Parser.h

@@ -960,10 +960,17 @@ class Parser {
   /// Parse a differentiation parameters clause.
   bool parseDifferentiationParametersClause(
       SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName);
+  
+  /// Parse a transposing parameters clause.
+  bool parseTransposingParametersClause(
+      SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName);
 
   /// Parse the @differentiating attribute.
   ParserResult<DifferentiatingAttr>
   parseDifferentiatingAttribute(SourceLoc AtLoc, SourceLoc Loc);
+  
+  ParserResult<TransposingAttr> parseTransposingAttribute(SourceLoc AtLoc,
+                                                          SourceLoc Loc);
 
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc);

include/swift/SIL/SILFunction.h

@@ -580,7 +580,7 @@ class SILFunction
   }
 
   /// Returns true if the function has parameters that are consumed by the
-  // callee.
+  /// callee.
   bool hasOwnedParameters() const {
     for (auto &ParamInfo : getLoweredFunctionType()->getParameters()) {
       if (ParamInfo.isConsumed())