autodiff builtins (#30624)

Define type signatures and SILGen for the following builtins:

```
/// Applies the {jvp|vjp} of `f` to `arg1`, ..., `argN`.
func applyDerivative_arityN_{jvp|vjp}(f, arg1, ..., argN) -> jvp/vjp return type

/// Applies the transpose of `f` to `arg`.
func applyTranspose_arityN(f, arg) -> transpose return type

/// Makes a differentiable function from the given `original`, `jvp`, and
/// `vjp` functions.
func differentiableFunction_arityN(original, jvp, vjp)

/// Makes a linear function from the given `original` and `transpose` functions.
func linearFunction_arityN(original, transpose)
```

Add SILGen FileCheck tests for all builtins.

Author: marcrasi

include/swift/AST/AutoDiff.h

@@ -441,6 +441,33 @@ GenericSignature getConstrainedDerivativeGenericSignature(
     GenericSignature derivativeGenSig, LookupConformanceFn lookupConformance,
     bool isTranspose = false);
 
+/// Retrieve config from the function name of a variant of
+/// `Builtin.applyDerivative`, e.g. `Builtin.applyDerivative_jvp_arity2`.
+/// Returns true if the function name is parsed successfully.
+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.
+/// `Builtin.differentiableFunction_arity1_throws`.
+/// Returns true if the function name is parsed successfully.
+bool getBuiltinDifferentiableOrLinearFunctionConfig(
+    StringRef operationName, unsigned &arity, bool &throws);
+
 } // end namespace autodiff
 
 } // end namespace swift

include/swift/AST/Builtins.def

@@ -469,6 +469,18 @@ BUILTIN_SIL_OPERATION(ConvertStrongToUnownedUnsafe, "convertStrongToUnownedUnsaf
 /// now.
 BUILTIN_SIL_OPERATION(ConvertUnownedUnsafeToGuaranteed, "convertUnownedUnsafeToGuaranteed", Special)
 
+/// applyDerivative
+BUILTIN_SIL_OPERATION(ApplyDerivative, "applyDerivative", Special)
+
+/// applyTranspose
+BUILTIN_SIL_OPERATION(ApplyTranspose, "applyTranspose", Special)
+
+/// differentiableFunction
+BUILTIN_SIL_OPERATION(DifferentiableFunction, "differentiableFunction", Special)
+
+/// linearFunction
+BUILTIN_SIL_OPERATION(LinearFunction, "linearFunction", Special)
+
 #undef BUILTIN_SIL_OPERATION
 
 // BUILTIN_RUNTIME_CALL - A call into a runtime function.

include/swift/AST/KnownIdentifiers.def

@@ -133,6 +133,8 @@ IDENTIFIER(withKeywordArguments)
 IDENTIFIER(wrapped)
 IDENTIFIER(wrappedValue)
 IDENTIFIER(wrapperValue)
+IDENTIFIER(differential)
+IDENTIFIER(pullback)
 
 // Kinds of layout constraints
 IDENTIFIER_WITH_NAME(UnknownLayout, "_UnknownLayout")

include/swift/AST/KnownProtocols.def

@@ -78,6 +78,7 @@ PROTOCOL_(SwiftNewtypeWrapper)
 PROTOCOL(CodingKey)
 PROTOCOL(Encodable)
 PROTOCOL(Decodable)
+PROTOCOL(AdditiveArithmetic)
 
 PROTOCOL_(ObjectiveCBridgeable)
 PROTOCOL_(DestructorSafeContainer)

include/swift/AST/Types.h

@@ -3353,6 +3353,8 @@ class AnyFunctionType : public TypeBase {
       IndexSubset *parameterIndices, AutoDiffLinearMapKind kind,
       LookupConformanceFn lookupConformance, bool makeSelfParamFirst = false);
 
+  AnyFunctionType *getWithoutDifferentiability() const;
+
   /// True if the parameter declaration it is attached to is guaranteed
   /// to not persist the closure for longer than the duration of the call.
   bool isNoEscape() const {

lib/AST/AutoDiff.cpp

@@ -279,6 +279,77 @@ GenericSignature autodiff::getConstrainedDerivativeGenericSignature(
       nullptr);
 }
 
+// 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)) {
+    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;
+  }
+}
+
+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();
+}
+
+bool autodiff::getBuiltinDifferentiableOrLinearFunctionConfig(
+    StringRef operationName, unsigned &arity, bool &throws) {
+  constexpr char differentiablePrefix[] = "differentiableFunction";
+  constexpr char linearPrefix[] = "linearFunction";
+  if (operationName.startswith(differentiablePrefix))
+    operationName = operationName.drop_front(sizeof(differentiablePrefix) - 1);
+  else if (operationName.startswith(linearPrefix))
+    operationName = operationName.drop_front(sizeof(linearPrefix) - 1);
+  else
+    return false;
+  parseAutoDiffBuiltinCommonConfig(operationName, arity, throws);
+  return operationName.empty();
+}
+
 Type TangentSpace::getType() const {
   switch (kind) {
   case Kind::TangentVector: