[AutoDiff] Clean up. (#27718)

- Store `AutoDiffConfig` in `SILDifferentiabilityWitness` instead of storing
  the individual components. This makes it cheaper to get an `AutoDiffConfig`.
- Unify parsing logic and diagnostics.
- Minor style changes.

Author: dan-zheng

docs/SIL.rst

@@ -5602,7 +5602,6 @@ Automatic Differentiation
 
 differentiable_function
 ```````````````````````
-
 ::
 
   sil-instruction ::= 'differentiable_function'
@@ -5638,7 +5637,6 @@ clause. In canonical SIL, a ``with_derivative`` clause is mandatory.
 
 linear_function
 ```````````````
-
 ::
 
   sil-instruction ::= 'linear_function'
@@ -5670,7 +5668,6 @@ In canonical SIL, a ``with`` clause is mandatory.
 
 differentiable_function_extract
 ```````````````````````````````
-
 ::
 
   sil-instruction ::= 'differentiable_function_extract'
@@ -5692,7 +5689,6 @@ Extracts the original function or a derivative function from the given
 
 linear_function_extract
 ```````````````````````
-
 ::
 
   sil-instruction ::= 'linear_function_extract'

include/swift/AST/AutoDiff.h

@@ -39,7 +39,7 @@ class SILFunctionType;
 typedef CanTypeWrapper<SILFunctionType> CanSILFunctionType;
 enum class SILLinkage : uint8_t;
 
-enum class DifferentiabilityKind: uint8_t {
+enum class DifferentiabilityKind : uint8_t {
   NonDifferentiable = 0,
   Normal = 1,
   Linear = 2
@@ -62,10 +62,10 @@ struct AutoDiffLinearMapKind {
 /// The kind of a derivative function.
 struct AutoDiffDerivativeFunctionKind {
   enum innerty : uint8_t {
-   // The Jacobian-vector products function.
-   JVP = 0,
-   // The vector-Jacobian products function.
-   VJP = 1
+    // The Jacobian-vector products function.
+    JVP = 0,
+    // The vector-Jacobian products function.
+    VJP = 1
   } rawValue;
 
   AutoDiffDerivativeFunctionKind() = default;
@@ -91,8 +91,8 @@ struct NormalDifferentiableFunctionTypeComponent {
       : rawValue(rawValue) {}
   NormalDifferentiableFunctionTypeComponent(
       AutoDiffDerivativeFunctionKind kind);
-  explicit NormalDifferentiableFunctionTypeComponent(unsigned rawValue) :
-      NormalDifferentiableFunctionTypeComponent((innerty)rawValue) {}
+  explicit NormalDifferentiableFunctionTypeComponent(unsigned rawValue)
+      : NormalDifferentiableFunctionTypeComponent((innerty)rawValue) {}
   explicit NormalDifferentiableFunctionTypeComponent(StringRef name);
   operator innerty() const { return rawValue; }
 
@@ -108,8 +108,8 @@ struct LinearDifferentiableFunctionTypeComponent {
   LinearDifferentiableFunctionTypeComponent() = default;
   LinearDifferentiableFunctionTypeComponent(innerty rawValue)
       : rawValue(rawValue) {}
-  explicit LinearDifferentiableFunctionTypeComponent(unsigned rawValue) :
-      LinearDifferentiableFunctionTypeComponent((innerty)rawValue) {}
+  explicit LinearDifferentiableFunctionTypeComponent(unsigned rawValue)
+      : LinearDifferentiableFunctionTypeComponent((innerty)rawValue) {}
   explicit LinearDifferentiableFunctionTypeComponent(StringRef name);
   operator innerty() const { return rawValue; }
 };
@@ -132,10 +132,10 @@ class ParsedAutoDiffParameter {
 
 public:
   ParsedAutoDiffParameter(SourceLoc loc, enum Kind kind, Value value)
-    : Loc(loc), Kind(kind), V(value) {}
+      : Loc(loc), Kind(kind), V(value) {}
 
   ParsedAutoDiffParameter(SourceLoc loc, enum Kind kind, unsigned index)
-  : Loc(loc), Kind(kind), V(index) {}
+      : Loc(loc), Kind(kind), V(index) {}
 
   static ParsedAutoDiffParameter getNamedParameter(SourceLoc loc,
                                                    Identifier name) {
@@ -251,6 +251,12 @@ struct AutoDiffConfig {
   IndexSubset *parameterIndices;
   IndexSubset *resultIndices;
   GenericSignature *derivativeGenericSignature;
+
+  /*implicit*/ AutoDiffConfig(IndexSubset *parameterIndices,
+                              IndexSubset *resultIndices,
+                              GenericSignature *derivativeGenericSignature)
+      : parameterIndices(parameterIndices), resultIndices(resultIndices),
+        derivativeGenericSignature(derivativeGenericSignature) {}
 };
 
 /// In conjunction with the original function declaration, identifies an

include/swift/AST/DiagnosticsParse.def

@@ -689,12 +689,6 @@ ERROR(sil_witness_protocol_conformance_not_found,none,
 // SIL differentiability witnesses
 ERROR(sil_diff_witness_expected_token,PointsToFirstBadToken,
       "expected '%0' in differentiability witness", (StringRef))
-ERROR(sil_diff_witness_expected_index_list,PointsToFirstBadToken,
-      "expected a space-separated list of indices, e.g. '0 1'", ())
-ERROR(sil_diff_witness_expected_parameter_index,PointsToFirstBadToken,
-      "expected a parameter index to differentiate with respect to", ())
-ERROR(sil_diff_witness_expected_result_index,PointsToFirstBadToken,
-      "expected a result index to differentiate with respect to", ())
 
 // SIL Coverage Map
 ERROR(sil_coverage_func_not_found, none,
@@ -1596,16 +1590,20 @@ ERROR(sil_attr_differentiable_expected_parameter_list,PointsToFirstBadToken,
       "expected an comma-separated list of parameter indices, e.g. (0, 1)", ())
 ERROR(sil_attr_differentiable_expected_rsquare,PointsToFirstBadToken,
       "expected ']' to end 'differentiable' attribute", ())
-ERROR(sil_attr_differentiable_expected_parameter_index,PointsToFirstBadToken,
-      "expected the index of a parameter to differentiate w.r.t.", ())
-ERROR(sil_attr_differentiable_expected_source_index,PointsToFirstBadToken,
-      "expected the index of a result to differentiate from", ())
 
 // SIL autodiff
-ERROR(sil_inst_autodiff_attr_expected_rsquare,PointsToFirstBadToken,
+ERROR(sil_autodiff_expected_lsquare,PointsToFirstBadToken,
+      "expected '[' to start the %0", (StringRef))
+ERROR(sil_autodiff_expected_rsquare,PointsToFirstBadToken,
       "expected ']' to complete the %0", (StringRef))
-ERROR(sil_inst_autodiff_expected_parameter_index,PointsToFirstBadToken,
+ERROR(sil_autodiff_expected_index_list,PointsToFirstBadToken,
+      "expected a space-separated list of indices, e.g. '0 1'", ())
+ERROR(sil_autodiff_expected_index_list_label,PointsToFirstBadToken,
+      "expected label '%0' in index list", (StringRef))
+ERROR(sil_autodiff_expected_parameter_index,PointsToFirstBadToken,
       "expected the index of a parameter to differentiate with respect to", ())
+ERROR(sil_autodiff_expected_result_index,PointsToFirstBadToken,
+      "expected the index of a result to differentiate from", ())
 ERROR(sil_inst_autodiff_operand_list_expected_lbrace,PointsToFirstBadToken,
       "expected '{' to start a derivative function list", ())
 ERROR(sil_inst_autodiff_operand_list_expected_comma,PointsToFirstBadToken,

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -48,12 +48,9 @@ class SILDifferentiabilityWitness
   SILLinkage linkage;
   /// The original function.
   SILFunction *originalFunction;
-  /// The parameter indices.
-  IndexSubset *parameterIndices;
-  /// The result indices.
-  IndexSubset *resultIndices;
-  /// The derivative generic signature (optional).
-  GenericSignature *derivativeGenericSignature;
+  /// The autodiff configuration: parameter indices, result indices, derivative
+  /// generic signature (optional).
+  AutoDiffConfig config;
   /// The JVP (Jacobian-vector products) derivative function.
   SILFunction *jvp;
   /// The VJP (vector-Jacobian products) derivative function.
@@ -75,9 +72,8 @@ class SILDifferentiabilityWitness
                               SILFunction *jvp, SILFunction *vjp,
                               bool isSerialized, DeclAttribute *attribute)
     : module(module), linkage(linkage), originalFunction(originalFunction),
-      parameterIndices(parameterIndices), resultIndices(resultIndices),
-      derivativeGenericSignature(derivativeGenSig), jvp(jvp), vjp(vjp),
-      serialized(isSerialized), attribute(attribute) {}
+      config(parameterIndices, resultIndices, derivativeGenSig), jvp(jvp),
+      vjp(vjp), serialized(isSerialized), attribute(attribute) {}
 
 public:
   static SILDifferentiabilityWitness *create(
@@ -90,14 +86,15 @@ class SILDifferentiabilityWitness
   SILModule &getModule() const { return module; }
   SILLinkage getLinkage() const { return linkage; }
   SILFunction *getOriginalFunction() const { return originalFunction; }
+  const AutoDiffConfig &getConfig() const { return config; }
   IndexSubset *getParameterIndices() const {
-    return parameterIndices;
+    return config.parameterIndices;
   }
   IndexSubset *getResultIndices() const {
-    return resultIndices;
+    return config.resultIndices;
   }
   GenericSignature *getDerivativeGenericSignature() const {
-    return derivativeGenericSignature;
+    return config.derivativeGenericSignature;
   }
   SILFunction *getJVP() const { return jvp; }
   SILFunction *getVJP() const { return vjp; }