[AutoDiff] Add 'Builtin.applyTranspose*'. (#28469)

Add support for applying the transpose function in a `@differentiable(linear)` function. The `Builtin.applyTranspose*` builtin takes a `@differentiable(linear)` function and a tangent vector, and returns the result of applying the transpose to the tangent vector.

Pseudo-declaration:
```swift
func applyTranspose_arity{arity}[_throws?{r}]<T...{arity}, R>(
    _: @differentiable (T...) {r}?throws -> R,
    _: R
) {r}?rethrows -> (T...)
    where T: Differentiable & AdditiveArithmetic, R: Differentiable & AdditiveArithmetic
```

This patch also renames `Builtin.autodiffApply` to `Builtin.applyDerivative` for clarity, and fixes a bug in `LinearDifferentiableSILFunctionTypeLowering` where it expected 3 component values in `rebuildAggregate` instead of 2.

Resolves SR-11844 and SR-11851.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -352,11 +352,24 @@ IndexSubset *getLoweredParameterIndices(IndexSubset *indices,
                                         AnyFunctionType *type);
 
 /// Retrieve config from the function name of a variant of
-/// `Builtin.autodiffApply`, e.g. `Builtin.autodiffApply_jvp_arity2`.
+/// `Builtin.applyDerivative`, e.g. `Builtin.applyDerivative_jvp_arity2`.
 /// Returns true if the function name is parsed successfully.
-bool getBuiltinAutoDiffApplyConfig(StringRef operationName,
-                                   AutoDiffDerivativeFunctionKind &kind,
-                                   unsigned &arity, bool &rethrows);
+bool getBuiltinApplyDerivativeConfig(
+    StringRef operationName, AutoDiffDerivativeFunctionKind &kind,
+    unsigned &arity, bool &rethrows);
+
+/// Retrieve config from the function name of a variant of
+/// `Builtin.applyTranspose`, e.g. `Builtin.applyTranspose_arity2`.
+/// Returns true if the function name is parsed successfully.
+bool getBuiltinApplyTransposeConfig(
+  StringRef operationName, unsigned &arity, bool &rethrows);
+
+/// Retrieve config from the function name of a variant of
+/// `Builtin.differentiableFunction` or `Builtin.linearFunction`, e.g.
+/// `Builtin.differentiableFunction_arity1_throws`.
+/// Returns true if the function name is parsed successfully.
+bool getBuiltinDifferentiableOrLinearFunctionConfig(
+    StringRef operationName, unsigned &arity, bool &throws);
 
 /// Retrieve config from the function name of a variant of
 /// `Builtin.differentiableFunction` or `Builtin.linearFunction`, e.g.

include/swift/AST/Builtins.def

@@ -445,8 +445,11 @@ BUILTIN_SIL_OPERATION(AllocWithTailElems, "allocWithTailElems", Special)
 BUILTIN_SIL_OPERATION(ProjectTailElems, "projectTailElems", Special)
 
 // SWIFT_ENABLE_TENSORFLOW
-/// autodiffApply
-BUILTIN_SIL_OPERATION(AutoDiffApply, "autodiffApply", Special)
+/// applyDerivative
+BUILTIN_SIL_OPERATION(ApplyDerivative, "applyDerivative", Special)
+
+/// applyTranspose
+BUILTIN_SIL_OPERATION(ApplyTranspose, "applyTranspose", Special)
 
 /// differentiableFunction
 BUILTIN_SIL_OPERATION(DifferentiableFunction, "differentiableFunction", Special)

lib/AST/AutoDiff.cpp

@@ -206,21 +206,12 @@ void autodiff::getSubsetParameterTypes(IndexSubset *subset,
   }
 }
 
-bool autodiff::getBuiltinAutoDiffApplyConfig(
-    StringRef operationName, AutoDiffDerivativeFunctionKind &kind,
-    unsigned &arity, bool &rethrows) {
-  constexpr char prefix[] = "autodiffApply";
-  if (!operationName.startswith(prefix))
-    return false;
-  operationName = operationName.drop_front(sizeof(prefix) - 1);
-  // Parse 'jvp' or 'vjp'.
-  constexpr char jvpPrefix[] = "_jvp";
-  constexpr char vjpPrefix[] = "_vjp";
-  if (operationName.startswith(jvpPrefix))
-    kind = AutoDiffDerivativeFunctionKind::JVP;
-  else if (operationName.startswith(vjpPrefix))
-    kind = AutoDiffDerivativeFunctionKind::VJP;
-  operationName = operationName.drop_front(sizeof(jvpPrefix) - 1);
+// Given the rest of a `Builtin.applyDerivative_{jvp|vjp}` or
+// `Builtin.applyTranspose` operation name, attempts to parse the arity and
+// throwing-ness from the operation name. Modifies the operation name argument
+// in place as substrings get dropped.
+static void parseAutoDiffBuiltinCommonConfig(
+    StringRef &operationName, unsigned &arity, bool &throws) {
   // Parse '_arity'.
   constexpr char arityPrefix[] = "_arity";
   if (operationName.startswith(arityPrefix)) {
@@ -233,14 +224,42 @@ bool autodiff::getBuiltinAutoDiffApplyConfig(
   } else {
     arity = 1;
   }
-  // Parse '_rethrows'.
-  constexpr char rethrowsPrefix[] = "_rethrows";
-  if (operationName.startswith(rethrowsPrefix)) {
-    operationName = operationName.drop_front(sizeof(rethrowsPrefix) - 1);
-    rethrows = true;
+  // Parse '_throws'.
+  constexpr char throwsPrefix[] = "_throws";
+  if (operationName.startswith(throwsPrefix)) {
+    operationName = operationName.drop_front(sizeof(throwsPrefix) - 1);
+    throws = true;
   } else {
-    rethrows = false;
+    throws = false;
   }
+}
+
+bool autodiff::getBuiltinApplyDerivativeConfig(
+    StringRef operationName, AutoDiffDerivativeFunctionKind &kind,
+    unsigned &arity, bool &throws) {
+  constexpr char prefix[] = "applyDerivative";
+  if (!operationName.startswith(prefix))
+    return false;
+  operationName = operationName.drop_front(sizeof(prefix) - 1);
+  // Parse 'jvp' or 'vjp'.
+  constexpr char jvpPrefix[] = "_jvp";
+  constexpr char vjpPrefix[] = "_vjp";
+  if (operationName.startswith(jvpPrefix))
+    kind = AutoDiffDerivativeFunctionKind::JVP;
+  else if (operationName.startswith(vjpPrefix))
+    kind = AutoDiffDerivativeFunctionKind::VJP;
+  operationName = operationName.drop_front(sizeof(jvpPrefix) - 1);
+  parseAutoDiffBuiltinCommonConfig(operationName, arity, throws);
+  return operationName.empty();
+}
+
+bool autodiff::getBuiltinApplyTransposeConfig(
+    StringRef operationName, unsigned &arity, bool &throws) {
+  constexpr char prefix[] = "applyTranspose";
+  if (!operationName.startswith(prefix))
+    return false;
+  operationName = operationName.drop_front(sizeof(prefix) - 1);
+  parseAutoDiffBuiltinCommonConfig(operationName, arity, throws);
   return operationName.empty();
 }
 
@@ -254,26 +273,7 @@ bool autodiff::getBuiltinDifferentiableOrLinearFunctionConfig(
     operationName = operationName.drop_front(sizeof(linearPrefix) - 1);
   else
     return false;
-  // Parse '_arity'.
-  constexpr char arityPrefix[] = "_arity";
-  if (operationName.startswith(arityPrefix)) {
-    operationName = operationName.drop_front(sizeof(arityPrefix) - 1);
-    auto arityStr = operationName.take_while(llvm::isDigit);
-    operationName = operationName.drop_front(arityStr.size());
-    auto converted = llvm::to_integer(arityStr, arity);
-    assert(converted); (void)converted;
-    assert(arity > 0);
-  } else {
-    arity = 1;
-  }
-  // Parse '_throws'.
-  constexpr char throwsPrefix[] = "_throws";
-  if (operationName.startswith(throwsPrefix)) {
-    operationName = operationName.drop_front(sizeof(throwsPrefix) - 1);
-    throws = true;
-  } else {
-    throws = false;
-  }
+  parseAutoDiffBuiltinCommonConfig(operationName, arity, throws);
   return operationName.empty();
 }
 

lib/AST/Builtins.cpp

@@ -985,7 +985,7 @@ static ValueDecl *getGetObjCTypeEncodingOperation(ASTContext &Context,
 // SWIFT_ENABLE_TENSORFLOW
 static ValueDecl *getAutoDiffApplyDerivativeFunction(
     ASTContext &Context, Identifier Id, AutoDiffDerivativeFunctionKind kind,
-    unsigned arity, bool rethrows) {
+    unsigned arity, bool throws) {
   assert(arity >= 1);
   // JVP:
   //   <...T...(arity), R> (@differentiable (...T) throws -> R, ...T)
@@ -1000,54 +1000,114 @@ static ValueDecl *getAutoDiffApplyDerivativeFunction(
   // Create type parameters and add conformance constraints.
   auto fnResultGen = makeGenericParam(arity);
   builder.addConformanceRequirement(fnResultGen, diffableProto);
-  SmallVector<decltype(fnResultGen), 2> fnArgGens;
+  SmallVector<decltype(fnResultGen), 2> fnParamGens;
   for (auto i : range(arity)) {
     auto T = makeGenericParam(i);
     builder.addConformanceRequirement(T, diffableProto);
-    fnArgGens.push_back(T);
+    fnParamGens.push_back(T);
   }
-  // Generator for the first argument, i.e. the @differentiable function.
+  // Generator for the first argument, i.e. the `@differentiable` function.
   BuiltinFunctionBuilder::LambdaGenerator firstArgGen {
     // Generator for the function type at the argument position, i.e. the
     // function being differentiated.
-    [=, &fnArgGens](BuiltinFunctionBuilder &builder) -> Type {
+    [=, &fnParamGens](BuiltinFunctionBuilder &builder) -> Type {
       FunctionType::ExtInfo ext;
       auto extInfo = FunctionType::ExtInfo()
           .withDifferentiabilityKind(DifferentiabilityKind::Normal)
-          .withNoEscape().withThrows(rethrows);
+          .withNoEscape().withThrows(throws);
       SmallVector<FunctionType::Param, 2> params;
-      for (auto &paramGen : fnArgGens)
+      for (auto &paramGen : fnParamGens)
         params.push_back(FunctionType::Param(paramGen.build(builder)));
       auto innerFunction = FunctionType::get(params,
                                              fnResultGen.build(builder));
       return innerFunction->withExtInfo(extInfo);
     }
   };
   // Eagerly build the type of the first arg, then use that to compute the type
-  // of the derivative function type.
-  auto *origFnTy =
+  // of the result.
+  auto *diffFnType =
       firstArgGen.build(builder)->castTo<AnyFunctionType>();
-  origFnTy = origFnTy->getWithoutDifferentiability()->withExtInfo(
-      origFnTy->getExtInfo().withNoEscape(false));
+  diffFnType = diffFnType->getWithoutDifferentiability()->withExtInfo(
+      diffFnType->getExtInfo().withNoEscape(false));
   auto *paramIndices = IndexSubset::get(
-      Context, SmallBitVector(origFnTy->getNumParams(), true));
+      Context, SmallBitVector(diffFnType->getNumParams(), true));
   // Generator for the resultant function type, i.e. the AD derivative function.
   BuiltinFunctionBuilder::LambdaGenerator resultGen{
       [=, &Context](BuiltinFunctionBuilder &builder) -> Type {
-        auto derivativeFnTy = origFnTy->getAutoDiffDerivativeFunctionType(
+        auto derivativeFnTy = diffFnType->getAutoDiffDerivativeFunctionType(
             paramIndices, /*resultIndex*/ 0, kind,
             LookUpConformanceInModule(Context.TheBuiltinModule));
         return derivativeFnTy->getResult();
       }};
   builder.addParameter(firstArgGen);
-  for (auto argGen : fnArgGens)
+  for (auto argGen : fnParamGens)
     builder.addParameter(argGen);
-  if (rethrows)
+  if (throws)
     builder.setRethrows();
   builder.setResult(resultGen);
   return builder.build(Id);
 }
 
+static ValueDecl *getAutoDiffApplyTransposeFunction(
+    ASTContext &Context, Identifier Id, unsigned arity, bool throws) {
+  assert(arity >= 1);
+  // <...T...(arity), R>
+  //     (@differentiable (...T) throws -> R, ...R.TangentVector)
+  //         rethrows -> (...T.TangentVector)
+  unsigned numGenericParams = 1 + arity;
+  BuiltinFunctionBuilder builder(Context, numGenericParams);
+  auto *diffableProto = Context.getProtocol(KnownProtocolKind::Differentiable);
+  auto *addArithProto =
+      Context.getProtocol(KnownProtocolKind::AdditiveArithmetic);
+  // Create type parameters and add conformance constraints.
+  auto linearFnResultGen = makeGenericParam(arity);
+  builder.addConformanceRequirement(linearFnResultGen, diffableProto);
+  builder.addConformanceRequirement(linearFnResultGen, addArithProto);
+  SmallVector<decltype(linearFnResultGen), 2> linearFnParamGens;
+  for (auto i : range(arity)) {
+    auto T = makeGenericParam(i);
+    builder.addConformanceRequirement(T, diffableProto);
+    builder.addConformanceRequirement(T, addArithProto);
+    linearFnParamGens.push_back(T);
+  }
+  // Generator for the first argument, i.e. the `@differentiable(linear)`
+  // function.
+  BuiltinFunctionBuilder::LambdaGenerator firstArgGen {
+    // Generator for the function type at the argument position, i.e. the
+    // function being differentiated.
+    [=, &linearFnParamGens](BuiltinFunctionBuilder &builder) -> Type {
+      FunctionType::ExtInfo ext;
+      auto extInfo = FunctionType::ExtInfo()
+          .withDifferentiabilityKind(DifferentiabilityKind::Linear)
+          .withNoEscape().withThrows(throws);
+      SmallVector<FunctionType::Param, 2> params;
+      for (auto &paramGen : linearFnParamGens)
+        params.push_back(FunctionType::Param(paramGen.build(builder)));
+      auto innerFunction = FunctionType::get(params,
+                                             linearFnResultGen.build(builder));
+      return innerFunction->withExtInfo(extInfo);
+    }
+  };
+  builder.addParameter(firstArgGen);
+  builder.addParameter(linearFnResultGen);
+  if (throws)
+    builder.setRethrows();
+  if (arity == 1)
+    builder.setResult(linearFnParamGens.front());
+  else {
+    BuiltinFunctionBuilder::LambdaGenerator tupleResultGen {
+      [&](BuiltinFunctionBuilder &builder) -> Type {
+        SmallVector<TupleTypeElt, 2> tupleElts;
+        for (auto linearFnParamGen : linearFnParamGens)
+          tupleElts.push_back(linearFnParamGen.build(builder));
+        return TupleType::get(tupleElts, Context);
+      }
+    };
+    builder.setResult(tupleResultGen);
+  }
+  return builder.build(Id);
+}
+
 static ValueDecl *getDifferentiableFunctionConstructor(
     ASTContext &Context, Identifier Id, unsigned arity, bool throws) {
   assert(arity >= 1);
@@ -1992,15 +2052,23 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getAllocWithTailElemsOperation(Context, Id, NumTailTypes);
   }
   // SWIFT_ENABLE_TENSORFLOW
-  if (OperationName.startswith("autodiffApply_")) {
+  if (OperationName.startswith("applyDerivative_")) {
     AutoDiffDerivativeFunctionKind kind;
     unsigned arity;
-    bool rethrows;
-    if (!autodiff::getBuiltinAutoDiffApplyConfig(OperationName, kind, arity,
-                                                 rethrows))
+    bool throws;
+    if (!autodiff::getBuiltinApplyDerivativeConfig(
+            OperationName, kind, arity, throws))
       return nullptr;
     return getAutoDiffApplyDerivativeFunction(Context, Id, kind, arity,
-                                              rethrows);
+                                              throws);
+  }
+  if (OperationName.startswith("applyTranspose_")) {
+    unsigned arity;
+    bool throws;
+    if (!autodiff::getBuiltinApplyTransposeConfig(
+            OperationName, arity, throws))
+      return nullptr;
+    return getAutoDiffApplyTransposeFunction(Context, Id, arity, throws);
   }
   if (OperationName.startswith("differentiableFunction_")) {
     unsigned arity;
@@ -2288,7 +2356,8 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getUnsafeGuaranteedEnd(Context, Id);
 
   // SWIFT_ENABLE_TENSORFLOW
-  case BuiltinValueKind::AutoDiffApply:
+  case BuiltinValueKind::ApplyDerivative:
+  case BuiltinValueKind::ApplyTranspose:
   case BuiltinValueKind::DifferentiableFunction:
   case BuiltinValueKind::LinearFunction:
     llvm_unreachable("Handled above");

lib/SIL/SILModule.cpp

@@ -310,8 +310,10 @@ const BuiltinInfo &SILModule::getBuiltinInfo(Identifier ID) {
   else if (OperationName.startswith("allocWithTailElems_"))
     Info.ID = BuiltinValueKind::AllocWithTailElems;
   // SWIFT_ENABLE_TENSORFLOW
-  else if (OperationName.startswith("autodiffApply_"))
-    Info.ID = BuiltinValueKind::AutoDiffApply;
+  else if (OperationName.startswith("applyDerivative_"))
+    Info.ID = BuiltinValueKind::ApplyDerivative;
+  else if (OperationName.startswith("applyTranspose_"))
+    Info.ID = BuiltinValueKind::ApplyTranspose;
   else if (OperationName.startswith("differentiableFunction_"))
     Info.ID = BuiltinValueKind::DifferentiableFunction;
   else if (OperationName.startswith("linearFunction_"))

lib/SIL/TypeLowering.cpp

@@ -962,7 +962,7 @@ namespace {
 
     SILValue rebuildAggregate(SILBuilder &B, SILLocation loc,
                               ArrayRef<SILValue> values) const override {
-      assert(values.size() == 3);
+      assert(values.size() == 2);
       auto fnTy = getLoweredType().castTo<SILFunctionType>();
       auto paramIndices = fnTy->getDifferentiationParameterIndices();
       return B.createLinearFunction(loc, paramIndices, values[0], values[1]);