[AutoDiff upstream] parsing for @differentiable function type (#27708)

Adds parsing for a type attribute `@differentiable`, which is optionally allowed to have argument `@differentiable(linear)`.

The typechecker currently rejects all uses of `@differentiable` with "error: attribute does not apply to type". Future work (https://bugs.swift.org/browse/TF-871 https://bugs.swift.org/browse/TF-873) will update the typechecker to allow this attribute in places where it is allowed.

Resolves https://bugs.swift.org/browse/TF-822.

Author: marcrasi

include/swift/AST/Attr.def

@@ -51,6 +51,7 @@ TYPE_ATTR(autoclosure)
 TYPE_ATTR(convention)
 TYPE_ATTR(noescape)
 TYPE_ATTR(escaping)
+TYPE_ATTR(differentiable)
 
 // SIL-specific attributes
 TYPE_ATTR(block_storage)

include/swift/AST/Attr.h

@@ -66,6 +66,10 @@ class TypeAttributes {
   Optional<StringRef> convention = None;
   Optional<StringRef> conventionWitnessMethodProtocol = None;
 
+  // Indicates whether the type's '@differentiable' attribute has a 'linear'
+  // argument.
+  bool linear = false;
+
   // For an opened existential type, the known ID.
   Optional<UUID> OpenedID;
 
@@ -80,7 +84,15 @@ class TypeAttributes {
   TypeAttributes() {}
 
   bool isValid() const { return AtLoc.isValid(); }
-  
+
+  bool isLinear() const {
+    assert(
+        !linear ||
+        (linear && has(TAK_differentiable)) &&
+            "Linear shouldn't have been true if there's no `@differentiable`");
+    return linear;
+  }
+
   void clearAttribute(TypeAttrKind A) {
     AttrLocs[A] = SourceLoc();
   }

include/swift/AST/DiagnosticsParse.def

@@ -1508,6 +1508,12 @@ ERROR(attr_specialize_expected_partial_or_full,none,
 ERROR(attr_implements_expected_member_name,PointsToFirstBadToken,
       "expected a member name as second parameter in '_implements' attribute", ())
 
+// differentiable
+ERROR(differentiable_attribute_expected_rparen,none,
+      "expected ')' in '@differentiable' attribute", ())
+ERROR(unexpected_argument_differentiable,none,
+      "unexpected argument '%0' in '@differentiable' attribute", (StringRef))
+
 //------------------------------------------------------------------------------
 // MARK: Generics parsing diagnostics
 //------------------------------------------------------------------------------

include/swift/AST/DiagnosticsSema.def

@@ -4728,6 +4728,13 @@ NOTE(previous_function_builder_here, none,
 ERROR(function_builder_arguments, none,
       "function builder attributes cannot have arguments", ())
 
+//------------------------------------------------------------------------------
+// MARK: differentiable programming diagnostics
+//------------------------------------------------------------------------------
+ERROR(experimental_differentiable_programming_disabled, none,
+      "differentiable programming is an experimental feature that is "
+      "currently disabled", ())
+
 #ifndef DIAG_NO_UNDEF
 # if defined(DIAG)
 #  undef DIAG

lib/AST/TypeRepr.cpp

@@ -299,6 +299,14 @@ void AttributedTypeRepr::printAttrs(ASTPrinter &Printer,
   if (hasAttr(TAK_escaping))
     Printer.printSimpleAttr("@escaping") << " ";
 
+  if (hasAttr(TAK_differentiable)) {
+    if (Attrs.isLinear()) {
+      Printer.printSimpleAttr("@differentiable(linear)") << " ";
+    } else {
+      Printer.printSimpleAttr("@differentiable") << " ";
+    }
+  }
+
   if (hasAttr(TAK_thin))
     Printer.printSimpleAttr("@thin") << " ";
   if (hasAttr(TAK_thick))

lib/Parse/ParseDecl.cpp

@@ -1944,6 +1944,53 @@ bool Parser::canParseTypeAttribute() {
                              /*justChecking*/ true);
 }
 
+/// Parses the '@differentiable' argument (no argument list, or '(linear)'),
+/// and sets the appropriate fields on `Attributes`.
+///
+/// \param emitDiagnostics - if false, doesn't emit diagnostics
+/// \returns true on error, false on success
+static bool parseDifferentiableAttributeArgument(Parser &P,
+                                                 TypeAttributes &Attributes,
+                                                 bool emitDiagnostics) {
+  Parser::BacktrackingScope backtrack(P);
+
+  // Match '( <identifier> )', and store the identifier token to `argument`.
+  if (!P.consumeIf(tok::l_paren))
+    return false;
+  auto argument = P.Tok;
+  if (!P.consumeIf(tok::identifier))
+    return false;
+  if (!P.consumeIf(tok::r_paren)) {
+    // Special case handling for '( <identifier> (' so that we don't produce the
+    // misleading diagnostic "expected ',' separator" when the real issue is
+    // that the user forgot the ')' closing the '@differentiable' argument list.
+    if (P.Tok.is(tok::l_paren)) {
+      backtrack.cancelBacktrack();
+      if (emitDiagnostics)
+        P.diagnose(P.Tok, diag::differentiable_attribute_expected_rparen);
+      return true;
+    }
+    return false;
+  }
+
+  // If the next token is an arrow, then the matched '( <identifier> )' is
+  // actually the parameter type list, not an argument to '@differentiable'.
+  if (P.Tok.is(tok::arrow))
+    return false;
+
+  backtrack.cancelBacktrack();
+
+  if (argument.getText() != "linear") {
+    if (emitDiagnostics)
+      P.diagnose(argument, diag::unexpected_argument_differentiable,
+                 argument.getText());
+    return true;
+  }
+
+  Attributes.linear = true;
+  return false;
+}
+
 /// \verbatim
 ///   attribute-type:
 ///     'noreturn'
@@ -2157,6 +2204,13 @@ bool Parser::parseTypeAttribute(TypeAttributes &Attributes, SourceLoc AtLoc,
     break;
   }
 
+  case TAK_differentiable: {
+    if (parseDifferentiableAttributeArgument(*this, Attributes,
+                                             /*emitDiagnostics=*/!justChecking))
+      return true;
+    break;
+  }
+
   // Convention attribute.
   case TAK_convention:
     Attributes.convention = conventionName;

lib/Sema/TypeCheckType.cpp

@@ -2062,7 +2062,7 @@ Type TypeResolver::resolveAttributedType(TypeAttributes &attrs,
   static const TypeAttrKind FunctionAttrs[] = {
     TAK_convention, TAK_pseudogeneric,
     TAK_callee_owned, TAK_callee_guaranteed, TAK_noescape, TAK_autoclosure,
-    TAK_escaping, TAK_yield_once, TAK_yield_many
+    TAK_differentiable, TAK_escaping, TAK_yield_once, TAK_yield_many
   };
 
   auto checkUnsupportedAttr = [&](TypeAttrKind attr) {
@@ -2211,6 +2211,12 @@ Type TypeResolver::resolveAttributedType(TypeAttributes &attrs,
         attrs.clearAttribute(TAK_autoclosure);
       }
 
+      if (attrs.has(TAK_differentiable) &&
+          !Context.LangOpts.EnableExperimentalDifferentiableProgramming) {
+        diagnose(attrs.getLoc(TAK_differentiable),
+                 diag::experimental_differentiable_programming_disabled);
+      }
+
       // Resolve the function type directly with these attributes.
       FunctionType::ExtInfo extInfo(rep, /*noescape=*/false,
                                     fnRepr->throws());

test/AutoDiff/Parse/differentiable_func_type.swift

@@ -0,0 +1,86 @@
+// RUN: %target-swift-frontend -dump-parse -verify %s | %FileCheck %s
+
+let a: @differentiable (Float) -> Float // okay
+// CHECK: (pattern_named 'a'
+// CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+let b: @differentiable(linear) (Float) -> Float // okay
+// CHECK: (pattern_named 'b'
+// CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+
+// Generic type test.
+struct A<T> {
+  func foo() {
+    let local: @differentiable(linear) (T) -> T // okay
+    // CHECK: (pattern_named 'local'
+    // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+  }
+}
+
+// expected-error @+1 {{expected ')' in '@differentiable' attribute}}
+let c: @differentiable(linear (Float) -> Float
+
+// expected-error @+1 {{expected ')' in '@differentiable' attribute}}
+let c: @differentiable(notValidArg (Float) -> Float
+
+// expected-error @+1 {{unexpected argument 'notValidArg' in '@differentiable' attribute}}
+let d: @differentiable(notValidArg) (Float) -> Float
+
+// Using 'linear' as a type
+struct B {
+  struct linear {}
+  let propertyB1: @differentiable (linear) -> Float // okay
+  // CHECK: (pattern_named 'propertyB1'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyB2: @differentiable(linear) (linear) -> linear // okay
+  // CHECK: (pattern_named 'propertyB2'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+
+  let propertyB3: @differentiable (linear, linear) -> linear // okay
+  // CHECK: (pattern_named 'propertyB3'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyB4: @differentiable (linear, Float) -> linear // okay
+  // CHECK: (pattern_named 'propertyB4'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyB5: @differentiable (Float, linear) -> linear // okay
+  // CHECK: (pattern_named 'propertyB5'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyB6: @differentiable(linear) (linear, linear, Float, linear)
+    -> Float // okay
+  // CHECK: (pattern_named 'propertyB6'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+
+  // expected-error @+1 {{expected ')' in '@differentiable' attribute}}
+  let propertyB7: @differentiable(linear (linear) -> Float
+}
+
+// Using 'linear' as a typealias
+struct C {
+  typealias linear = (C) -> C
+  let propertyC1: @differentiable (linear) -> Float // okay
+  // CHECK: (pattern_named 'propertyC1'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyC2: @differentiable(linear) (linear) -> linear // okay
+  // CHECK: (pattern_named 'propertyC2'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+
+  let propertyC3: @differentiable linear // okay
+  // CHECK: (pattern_named 'propertyC3'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable{{[^(]}}
+
+  let propertyC4: linear // okay
+  // CHECK: (pattern_named 'propertyC4'
+
+  let propertyC5: @differentiable(linear) linear // okay
+  // CHECK: (pattern_named 'propertyC5'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+
+  let propertyC6: @differentiable(linear) @convention(c) linear // okay
+  // CHECK: (pattern_named 'propertyC6'
+  // CHECK-NEXT: (type_attributed attrs=@differentiable(linear)
+}

test/AutoDiff/Sema/differentiable_features_disabled.swift

@@ -0,0 +1,4 @@
+// RUN: %target-swift-frontend -typecheck -verify %s
+
+// expected-error @+1 {{differentiable programming is an experimental feature that is currently disabled}}
+let _: @differentiable (Float) -> Float

test/AutoDiff/Sema/differentiable_func_type.swift

@@ -0,0 +1,6 @@
+// RUN: %target-swift-frontend -enable-experimental-differentiable-programming -typecheck -verify %s
+
+// expected-error @+1 {{@differentiable attribute only applies to function types}}
+let _: @differentiable Float
+
+let _: @differentiable (Float) -> Float