[AutoDiff] Clean up.

- 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

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; }

lib/ParseSIL/ParseSIL.cpp

@@ -992,7 +992,7 @@ static bool parseDifferentiableAttr(
   if (P.parseSpecificIdentifier(
           "source", diag::sil_attr_differentiable_expected_keyword, "source") ||
       P.parseUnsignedInteger(SourceIndex, LastLoc,
-           diag::sil_attr_differentiable_expected_source_index))
+           diag::sil_autodiff_expected_result_index))
     return true;
   // Parse 'wrt'.
   if (P.parseSpecificIdentifier(
@@ -1005,7 +1005,7 @@ static bool parseDifferentiableAttr(
     unsigned Index;
     // TODO: Reject non-ascending parameter index lists.
     if (P.parseUnsignedInteger(Index, LastLoc,
-            diag::sil_attr_differentiable_expected_parameter_index))
+            diag::sil_autodiff_expected_parameter_index))
       return true;
     ParamIndices.push_back(Index);
     return false;
@@ -2939,12 +2939,11 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
       while (P.Tok.is(tok::integer_literal)) {
         unsigned index;
         if (P.parseUnsignedInteger(index, lastLoc,
-              diag::sil_inst_autodiff_expected_parameter_index))
+              diag::sil_autodiff_expected_parameter_index))
           return true;
         parameterIndices.push_back(index);
       }
-      if (P.parseToken(tok::r_square,
-                       diag::sil_inst_autodiff_attr_expected_rsquare,
+      if (P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
                        "parameter index list"))
         return true;
     }
@@ -3002,12 +3001,11 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
       while (P.Tok.is(tok::integer_literal)) {
         unsigned index;
         if (P.parseUnsignedInteger(index, lastLoc,
-              diag::sil_inst_autodiff_expected_parameter_index))
+              diag::sil_autodiff_expected_parameter_index))
           return true;
         parameterIndices.push_back(index);
       }
-      if (P.parseToken(tok::r_square,
-                       diag::sil_inst_autodiff_attr_expected_rsquare,
+      if (P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
                        "parameter index list"))
         return true;
     }
@@ -3050,8 +3048,7 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
             diag::sil_inst_autodiff_expected_differentiable_extractee_kind) ||
         parseSILIdentifierSwitch(extractee, extracteeNames,
             diag::sil_inst_autodiff_expected_differentiable_extractee_kind) ||
-        P.parseToken(tok::r_square,
-                     diag::sil_inst_autodiff_attr_expected_rsquare,
+        P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
                      "extractee kind"))
       return true;
     if (parseTypedValueRef(functionOperand, B) ||
@@ -3073,8 +3070,7 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
             diag::sil_inst_autodiff_expected_linear_extractee_kind) ||
         parseSILIdentifierSwitch(extractee, extracteeNames,
             diag::sil_inst_autodiff_expected_linear_extractee_kind) ||
-        P.parseToken(tok::r_square,
-                     diag::sil_inst_autodiff_attr_expected_rsquare,
+        P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
                      "extractee kind"))
       return true;
     if (parseTypedValueRef(functionOperand, B) ||
@@ -6941,47 +6937,35 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   };
 
   SourceLoc lastLoc = P.getEndOfPreviousLoc();
-  // Parse an index subset, prefaced with the given label.
-  auto parseIndexSubset =
-      [&](StringRef label, IndexSubset *& indexSubset) -> bool {
-    if (P.parseToken(tok::l_square, diag::sil_diff_witness_expected_token, "["))
-      return true;
-    if (P.parseSpecificIdentifier(
-          label, diag::sil_diff_witness_expected_token, label))
-      return true;
-    // Parse parameter index list.
-    SmallVector<unsigned, 8> paramIndices;
-    // Function that parses an index into `paramIndices`. Returns true on error.
-    auto parseParam = [&]() -> bool {
+  // Parse an index set, prefaced with the given label.
+  auto parseIndexSet = [&](StringRef label, SmallVectorImpl<unsigned> &indices,
+                           const Diagnostic &parseIndexDiag) -> bool {
+    // Parse `[<label> <integer_literal>...]`.
+    if (P.parseToken(tok::l_square, diag::sil_autodiff_expected_lsquare,
+                     "index list") ||
+        P.parseSpecificIdentifier(
+            label, diag::sil_autodiff_expected_index_list_label, label))
+      return true;
+    while (P.Tok.is(tok::integer_literal)) {
       unsigned index;
-      // TODO: Reject non-ascending index lists.
-      if (P.parseUnsignedInteger(index, lastLoc,
-              diag::sil_diff_witness_expected_index_list))
+      if (P.parseUnsignedInteger(index, lastLoc, parseIndexDiag))
         return true;
-      paramIndices.push_back(index);
-      return false;
-    };
-    // Parse first.
-    if (parseParam())
-      return true;
-    // Parse rest.
-    while (P.Tok.isNot(tok::r_square))
-      if (parseParam())
-        return true;
-    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_token, "]"))
+      indices.push_back(index);
+    }
+    if (P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
+                     "index list"))
       return true;
-    auto maxIndexRef =
-        std::max_element(paramIndices.begin(), paramIndices.end());
-    indexSubset = IndexSubset::get(
-        P.Context, maxIndexRef ? *maxIndexRef + 1 : 0, paramIndices);
     return false;
   };
   // Parse parameter and result indices.
-  IndexSubset *parameterIndices = nullptr;
-  IndexSubset *resultIndices = nullptr;
-  if (parseIndexSubset("parameters", parameterIndices))
+  SmallVector<unsigned, 8> parameterIndices;
+  SmallVector<unsigned, 8> resultIndices;
+  // Parse parameter and result indices.
+  if (parseIndexSet("parameters", parameterIndices,
+                    diag::sil_autodiff_expected_parameter_index))
     return true;
-  if (parseIndexSubset("results", resultIndices))
+  if (parseIndexSet("results", resultIndices,
+                    diag::sil_autodiff_expected_result_index))
     return true;
 
   // Parse a trailing 'where' clause (optional).
@@ -6995,7 +6979,8 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     P.parseGenericWhereClause(whereLoc, derivativeRequirementReprs,
                               firstTypeInComplete,
                               /*AllowLayoutConstraints*/ false);
-    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_token, "]"))
+    if (P.parseToken(tok::r_square, diag::sil_autodiff_expected_rsquare,
+                     "where clause"))
       return true;
   }
 
@@ -7053,8 +7038,13 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
       return true;
   }
 
+  auto origFnType = originalFn->getLoweredFunctionType();
+  auto *parameterIndexSet = IndexSubset::get(
+      P.Context, origFnType->getNumParameters(), parameterIndices);
+  auto *resultIndexSet = IndexSubset::get(
+      P.Context, origFnType->getNumResults(), resultIndices);
   SILDifferentiabilityWitness::create(
-      M, *linkage, originalFn, parameterIndices, resultIndices,
+      M, *linkage, originalFn, parameterIndexSet, resultIndexSet,
       derivativeGenSig, jvp, vjp, isSerialized);
   return false;
 }

lib/SIL/SILDifferentiabilityWitness.cpp

@@ -34,7 +34,5 @@ SILDifferentiabilityWitness *SILDifferentiabilityWitness::create(
 }
 
 SILDifferentiabilityWitnessKey SILDifferentiabilityWitness::getKey() const {
-  AutoDiffConfig config{parameterIndices, resultIndices,
-                        derivativeGenericSignature};
-  return std::make_pair(originalFunction->getName(), config);
+  return std::make_pair(originalFunction->getName(), getConfig());
 }

lib/SILGen/SILGen.cpp

@@ -773,8 +773,8 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
       if (auto *vjpDecl = diffAttr->getVJPFunction())
         vjp = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
       auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
-      AutoDiffConfig config{diffAttr->getParameterIndices(), resultIndices,
-                            diffAttr->getDerivativeGenericSignature()};
+      AutoDiffConfig config(diffAttr->getParameterIndices(), resultIndices,
+                            diffAttr->getDerivativeGenericSignature());
       emitDifferentiabilityWitness(AFD, F, config, jvp, vjp);
     }
   }