[stdlib] [Sema] Add broadcasting 'adding(_:)' and 'subtracting(_:)' to 'VectorProtocol'. (#25525)

- Add broadcasting 'adding(_:)' and 'subtracting(_:)' to 'VectorProtocol'. These are important for aggregate algorithms such as machine learning optimizers (tensorflow/swift-apis#218).
- Make their implementations compiler-derivable.
- Add operators for `adding(_:)` and `subtracting(_:)` in a protocol extension.
- Comment out all operators in protocol extensions for `VectorProtocol` because a source breakage has been found.

Author: rxwei

include/swift/AST/KnownIdentifiers.def

@@ -135,7 +135,12 @@ IDENTIFIER_(typeList)
 // AdditiveArithmetic, VectorProtocol
 IDENTIFIER(zero)
 IDENTIFIER(VectorSpaceScalar)
+IDENTIFIER(adding)
+IDENTIFIER(subtracting)
 IDENTIFIER(scaled)
+IDENTIFIER(by)
+IDENTIFIER(scale)
+IDENTIFIER(x)
 // Differentiable
 IDENTIFIER(AllDifferentiableVariables)
 IDENTIFIER(TangentVector)

lib/Sema/DerivedConformanceVectorProtocol.cpp

@@ -202,30 +202,23 @@ static void deriveBodyVectorProtocol_method(AbstractFunctionDecl *funcDecl,
       BraceStmt::create(C, SourceLoc(), returnStmt, SourceLoc(), true));
 }
 
-// Synthesize body for `scaled(by:)`.
-static void deriveBodyVectorProtocol_scalarMultiply(
-    AbstractFunctionDecl *funcDecl, void *) {
-  auto &C = funcDecl->getASTContext();
-  deriveBodyVectorProtocol_method(funcDecl, C.Id_scaled, C.getIdentifier("by"));
-}
-
 // Synthesize function declaration for a `VectorProtocol` method requirement.
 static ValueDecl *deriveVectorProtocol_method(
-    DerivedConformance &derived, Identifier methodName, Identifier argumentName,
-    Identifier parameterName, Type parameterType, Type returnType,
-    AbstractFunctionDecl::BodySynthesizer bodySynthesizer) {
+    DerivedConformance &derived, Identifier methodBaseName,
+    Identifier argumentLabel, Identifier parameterName, Type parameterType,
+    Type returnType, AbstractFunctionDecl::BodySynthesizer bodySynthesizer) {
   auto nominal = derived.Nominal;
   auto &TC = derived.TC;
   auto &C = derived.TC.Context;
   auto parentDC = derived.getConformanceContext();
 
   auto *param =
       new (C) ParamDecl(VarDecl::Specifier::Default, SourceLoc(), SourceLoc(),
-                        argumentName, SourceLoc(), parameterName, parentDC);
+                        argumentLabel, SourceLoc(), parameterName, parentDC);
   param->setInterfaceType(parameterType);
   ParameterList *params = ParameterList::create(C, {param});
 
-  DeclName declName(C, methodName, params);
+  DeclName declName(C, methodBaseName, params);
   auto funcDecl = FuncDecl::create(C, SourceLoc(), StaticSpellingKind::None,
                                    SourceLoc(), declName, SourceLoc(),
                                    /*Throws*/ false, SourceLoc(),
@@ -258,8 +251,13 @@ static ValueDecl *deriveVectorProtocol_method(
   return funcDecl;
 }
 
-// Synthesize the `scaled(by:)` function declaration.
-static ValueDecl *deriveVectorProtocol_scaled(DerivedConformance &derived) {
+/// Synthesize a method declaration that has the following signture:
+///   func {methodBaseName}(
+///     {argumentLabel} {parameterName}: VectorSpaceScalar
+///   ) -> Self
+static ValueDecl *deriveVectorProtocol_unaryMethodOnScalar(
+    DerivedConformance &derived, Identifier methodBaseName,
+    Identifier argumentLabel, Identifier parameterName) {
   auto &C = derived.TC.Context;
   auto *nominal = derived.Nominal;
   auto *parentDC = derived.getConformanceContext();
@@ -268,18 +266,32 @@ static ValueDecl *deriveVectorProtocol_scaled(DerivedConformance &derived) {
   auto scalarType = deriveVectorProtocol_VectorSpaceScalar(nominal, parentDC)
       ->mapTypeOutOfContext();
 
+  auto bodySynthesizer = [](AbstractFunctionDecl *funcDecl, void *ctx) {
+    auto methodNameAndLabel = reinterpret_cast<Identifier *>(ctx);
+    deriveBodyVectorProtocol_method(
+        funcDecl, methodNameAndLabel[0], methodNameAndLabel[1]);
+  };
+  Identifier baseNameAndLabel[2] = {methodBaseName, argumentLabel};
   return deriveVectorProtocol_method(
-      derived, C.Id_scaled, C.getIdentifier("by"), C.getIdentifier("scalar"),
-      scalarType, selfInterfaceType,
-      {deriveBodyVectorProtocol_scalarMultiply, nullptr});
+      derived, methodBaseName, argumentLabel, parameterName, scalarType,
+      selfInterfaceType,
+      {bodySynthesizer, C.AllocateCopy(baseNameAndLabel).data()});
 }
 
 ValueDecl *DerivedConformance::deriveVectorProtocol(ValueDecl *requirement) {
   // Diagnose conformances in disallowed contexts.
   if (checkAndDiagnoseDisallowedContext(requirement))
     return nullptr;
+  auto &C = requirement->getASTContext();
   if (requirement->getBaseName() == TC.Context.Id_scaled)
-    return deriveVectorProtocol_scaled(*this);
+    return deriveVectorProtocol_unaryMethodOnScalar(
+        *this, C.Id_scaled, C.Id_by, C.Id_scale);
+  if (requirement->getBaseName() == TC.Context.Id_adding)
+    return deriveVectorProtocol_unaryMethodOnScalar(
+        *this, C.Id_adding, Identifier(), C.Id_x);
+  if (requirement->getBaseName() == TC.Context.Id_subtracting)
+    return deriveVectorProtocol_unaryMethodOnScalar(
+        *this, C.Id_subtracting, Identifier(), C.Id_x);
   TC.diagnose(requirement->getLoc(), diag::broken_vector_protocol_requirement);
   return nullptr;
 }

lib/Sema/DerivedConformances.cpp

@@ -346,6 +346,17 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,
         return getRequirement(KnownProtocolKind::VectorProtocol);
     }
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // VectorProtocol.adding(_:)
+    // VectorProtocol.subtracting(_:)
+    if (name.isCompoundName() &&
+        (name.getBaseName() == ctx.Id_adding ||
+         name.getBaseName() == ctx.Id_subtracting)) {
+      auto argumentNames = name.getArgumentNames();
+      if (argumentNames.size() == 1 && argumentNames[0].empty())
+        return getRequirement(KnownProtocolKind::VectorProtocol);
+    }
+
     // SWIFT_ENABLE_TENSORFLOW
     // TensorArrayProtocol._unpackTensorHandles(into:)
     if (name.isCompoundName() && 

stdlib/public/core/AutoDiff.swift

@@ -26,6 +26,14 @@ public protocol VectorProtocol : AdditiveArithmetic {
   /// The type of scalars in the vector space.
   associatedtype VectorSpaceScalar : AdditiveArithmetic
 
+  func adding(_ x: VectorSpaceScalar) -> Self
+
+  mutating func add(_ x: VectorSpaceScalar)
+
+  func subtracting(_ x: VectorSpaceScalar) -> Self
+
+  mutating func subtract(_ x: VectorSpaceScalar)
+
   /// Returns `self` multiplied by the given scalar.
   func scaled(by scalar: VectorSpaceScalar) -> Self
 
@@ -34,28 +42,66 @@ public protocol VectorProtocol : AdditiveArithmetic {
 }
 
 public extension VectorProtocol {
+  mutating func add(_ x: VectorSpaceScalar) {
+    self = adding(x)
+  }
+
+  mutating func subtract(_ x: VectorSpaceScalar) {
+    self = subtracting(x)
+  }
+
   mutating func scale(by scalar: VectorSpaceScalar) {
     self = scaled(by: scalar)
   }
+}
+
+/* Note: These default-implemented opreators will slow down type-checking
+   performance and break existing code.
+
+public extension VectorProtocol {
+  static func + (lhs: Self, rhs: VectorSpaceScalar) -> Self {
+    lhs.adding(rhs)
+  }
+
+  static func + (lhs: VectorSpaceScalar, rhs: Self) -> Self {
+    rhs.adding(lhs)
+  }
+
+  static func += (lhs: inout Self, rhs: VectorSpaceScalar) {
+    lhs.add(rhs)
+  }
+
+  static func - (lhs: Self, rhs: VectorSpaceScalar) -> Self {
+    lhs.subtracting(rhs)
+  }
+
+  static func -= (lhs: inout Self, rhs: VectorSpaceScalar) {
+    lhs.subtract(rhs)
+  }
 
   static func * (lhs: Self, rhs: VectorSpaceScalar) -> Self {
-    return lhs.scaled(by: rhs)
+    lhs.scaled(by: rhs)
   }
 
   static func * (lhs: VectorSpaceScalar, rhs: Self) -> Self {
-    return rhs.scaled(by: lhs)
+    rhs.scaled(by: lhs)
   }
 
   static func *= (lhs: inout Self, rhs: VectorSpaceScalar) {
     lhs.scale(by: rhs)
   }
 }
 
-public extension VectorProtocol where VectorSpaceScalar: SignedNumeric {
+public extension VectorProtocol where VectorSpaceScalar : SignedNumeric {
+  static func - (lhs: VectorSpaceScalar, rhs: Self) -> Self {
+    -rhs.adding(lhs)
+  }
+
   static prefix func - (x: Self) -> Self {
     .zero - x
   }
 }
+*/
 
 /// A type that mathematically represents a differentiable manifold whose
 /// tangent spaces are finite-dimensional.

stdlib/public/core/FloatingPointTypes.swift.gyb

@@ -1879,8 +1879,28 @@ extension ${Self} : Strideable {
 extension ${Self} : VectorProtocol {
   public typealias VectorSpaceScalar = ${Self}
 
+  public func adding(_ x: ${Self}) -> ${Self} {
+    self + x
+  }
+
+  public mutating func add(_ x: ${Self}) {
+    self += x
+  }
+
+  public func subtracting(_ x: ${Self}) -> ${Self} {
+    self - x
+  }
+
+  public mutating func subtract(_ x: ${Self}) {
+    self -= x
+  }
+
   public func scaled(by scalar: ${Self}) -> ${Self} {
-    return self * scalar
+    self * scalar
+  }
+
+  public mutating func scale(by scalar: ${Self}) {
+    self *= scalar
   }
 }
 

test/AutoDiff/refcounting.swift

@@ -23,6 +23,8 @@ public struct Vector : AdditiveArithmetic, VectorProtocol, Differentiable, Equat
   @differentiable(vjp: fakeVJP)
   public static func - (lhs: Vector, rhs: Vector) -> Vector { abort() }
 
+  public func adding(_ scalar: Float) -> Vector { abort() }
+  public func subtracting(_ scalar: Float) -> Vector { abort() }
   public func scaled(by scalar: Float) -> Vector { abort() }
 
   public static func fakeVJP(lhs: Vector, rhs: Vector) -> (Vector, (Vector) -> (Vector, Vector)) { abort() }

test/Sema/struct_key_path_iterable.swift

@@ -11,6 +11,8 @@ extension Tensor : Equatable where Scalar : Equatable {}
 extension Tensor : AdditiveArithmetic where Scalar : AdditiveArithmetic {}
 extension Tensor : VectorProtocol where Scalar : AdditiveArithmetic {
   typealias VectorSpaceScalar = Scalar
+  func adding(_: Scalar) -> Self { self }
+  func subtracting(_: Scalar) -> Self { self }
   func scaled(by scalar: Scalar) -> Self { self }
 }
 
@@ -55,6 +57,12 @@ extension TensorParameters : VectorProtocol {
     return TensorParameters(w: lhs.w + rhs.w, b: lhs.b + rhs.b)
   }
   typealias VectorSpaceScalar = Float
+  func adding(_ x: VectorSpaceScalar) -> TensorParameters {
+    return TensorParameters(w: w.adding(x), b: b.adding(x))
+  }
+  func subtracting(_ x: VectorSpaceScalar) -> TensorParameters {
+    return TensorParameters(w: w.subtracting(x), b: b.subtracting(x))
+  }
   func scaled(by scalar: VectorSpaceScalar) -> TensorParameters {
     return TensorParameters(w: w.scaled(by: scalar), b: b.scaled(by: scalar))
   }
@@ -100,8 +108,8 @@ struct DummyOptimizer<P : KeyPathIterable, Scalar : BinaryFloatingPoint>
     parameters: inout P, withGradients gradients: P
   ) {
     for kp in parameters.recursivelyAllWritableKeyPaths(to: Tensor<Scalar>.self) {
-      firstMoments[keyPath: kp] *= learningRate
-      parameters[keyPath: kp] -= learningRate * parameters[keyPath: kp]
+      firstMoments[keyPath: kp].scale(by: learningRate)
+      parameters[keyPath: kp] -= parameters[keyPath: kp].scaled(by: learningRate)
     }
   }
 }

test/Sema/struct_vector_protocol.swift

@@ -7,11 +7,28 @@ _ x: inout T, scalar: T.VectorSpaceScalar
   // Test `AdditiveArithmetic` requirements: `zero`, `+`, `-`.
   let zero = T.zero
   x += x + zero
-  x -= x - zero
-  // Test `VectorProtocol` requirements: `VectorSpaceScalar`, `*`.
-  x *= scalar
-  _ = scalar * x
-  _ = x * scalar
+  x += x - zero
+  // Test `VectorProtocol` requirements: `VectorSpaceScalar`, `adding(_:)`, `add(_:)`
+  // `subtracting(_:)`, `subtract(_:)`, `scaled(by:)`, and `scale(by:)`.
+  x.add(scalar)
+  x.add(scalar)
+  x.scale(by: scalar)
+  _ = x.adding(scalar)
+  _ = x.subtracting(scalar)
+  _ = x.scaled(by: scalar)
+   
+  // NOTE: Operators have been disabled for type checker performance reasons.
+  // x += x + zero
+  // x -= x - zero
+  // Test `VectorProtocol` requirements: `VectorSpaceScalar`, `+`, `-`, `*`.
+  // x += scalar
+  // x -= scalar
+  // x *= scalar
+  // _ = x + scalar
+  // _ = scalar + x
+  // _ = x - scalar
+  // _ = scalar * x
+  // _ = x * scalar
 }
 
 struct Float2: VectorProtocol {

utils/update_checkout/update-checkout-config.json

@@ -375,7 +375,7 @@
                 "clang-tools-extra": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",
                 "libcxx": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",
                 "tensorflow": "ebc41609e27dcf0998d8970e77a2e1f53e13ac86",
-                "tensorflow-swift-apis": "5d3ef57b501781f1bab3b4ca85e8b8fc91671c14",
+                "tensorflow-swift-apis": "7a3ed481bba53a7cd82f8a46c0df9f09a6e9747f",
                 "indexstore-db": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",
                 "sourcekit-lsp": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a"
             }