Update differentiability witness syntax.

Print original function name in comment.

```
// differentiability witness for foo
sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo
```

Author: dan-zheng

lib/ParseSIL/ParseSIL.cpp

@@ -6805,17 +6805,17 @@ static void convertRequirements(Parser &P, SILFunction *F,
 
 /// decl-sil-differentiability-witness ::=
 ///   'sil_differentiability_witness'
+///   '[' 'parameters' index-subset ']'
+///   '[' 'results' index-subset ']'
+///   ('[' 'where' derivatve-generic-signature-requirements ']')?
 ///   sil-function-name ':' sil-type
-///   'parameters' autodiff-index-subset
-///   'results' autodiff-index-subset
-///   ('where' generic-signature)?
 ///   '{'
 ///   ('jvp' sil-function-name ':' sil-type)?
 ///   ('vjp' sil-function-name ':' sil-type)?
 ///   '}'
 ///
-/// autodiff-index-subset ::=
-///   '(' [0-9]+ (',', [0-9]+)* ')'
+/// index-subset ::=
+///   [0-9]+ (' ' [0-9]+)*
 bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   P.consumeToken(tok::kw_sil_differentiability_witness);
   SILParser State(P);
@@ -6853,21 +6853,17 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     State.TUState.PotentialZombieFns.insert(fn);
     return false;
   };
-  // Parse original function name.
-  SILFunction *originalFn;
-  if (parseFunctionName(originalFn))
-    return true;
 
   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_keyword,
+                     "["))
+      return true;
     if (P.parseSpecificIdentifier(
           label, diag::sil_diff_witness_expected_keyword, label))
       return true;
-    if (P.parseToken(tok::l_paren, diag::sil_diff_witness_expected_keyword,
-                     "("))
-      return true;
     // Parse parameter index list.
     SmallVector<unsigned, 8> paramIndices;
     // Function that parses an index into `paramIndices`. Returns true on error.
@@ -6884,11 +6880,11 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     if (parseParam())
       return true;
     // Parse rest.
-    while (P.consumeIf(tok::comma))
+    while (P.Tok.isNot(tok::r_square))
       if (parseParam())
         return true;
-    if (P.parseToken(tok::r_paren, diag::sil_diff_witness_expected_keyword,
-                     "("))
+    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_keyword,
+                     "]"))
       return true;
     auto maxIndexRef =
         std::max_element(paramIndices.begin(), paramIndices.end());
@@ -6907,18 +6903,33 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   // Parse a trailing 'where' clause (optional).
   // This represents derivative generic signature requirements.
   GenericSignature *derivativeGenSig = nullptr;
-  if (P.Tok.is(tok::kw_where)) {
-    SourceLoc whereLoc;
-    SmallVector<RequirementRepr, 4> requirementReprs;
+  SourceLoc whereLoc;
+  SmallVector<RequirementRepr, 4> derivativeRequirementReprs;
+  if (P.Tok.is(tok::l_square) && P.peekToken().is(tok::kw_where)) {
+    P.consumeToken(tok::l_square);
     bool firstTypeInComplete;
-    P.parseGenericWhereClause(whereLoc, requirementReprs, firstTypeInComplete,
+    P.parseGenericWhereClause(whereLoc, derivativeRequirementReprs,
+                              firstTypeInComplete,
                               /*AllowLayoutConstraints*/ false);
-    auto *whereClause = TrailingWhereClause::create(
-        originalFn->getModule().getASTContext(), whereLoc, requirementReprs);
+    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_keyword,
+                     "]"))
+      return true;
+  }
+
+  // Parse original function name.
+  SILFunction *originalFn;
+  if (parseFunctionName(originalFn))
+    return true;
+
+  // Resolve derivative requirements.
+  if (!derivativeRequirementReprs.empty()) {
     SmallVector<Requirement, 4> requirements;
+    auto *whereClause = TrailingWhereClause::create(
+        originalFn->getModule().getASTContext(), whereLoc,
+        derivativeRequirementReprs);
     convertRequirements(P, originalFn, whereClause->getRequirements(),
                         requirements);
-    assert(requirements.size() == requirementReprs.size());
+    assert(requirements.size() == derivativeRequirementReprs.size());
     derivativeGenSig = evaluateOrDefault(
         P.Context.evaluator,
         AbstractGenericSignatureRequest{

lib/SIL/SILPrinter.cpp

@@ -3060,24 +3060,24 @@ void SILDefaultWitnessTable::dump() const {
 // SWIFT_ENABLE_TENSORFLOW
 void SILDifferentiabilityWitness::print(
     llvm::raw_ostream &OS, bool verbose) const {
+  OS << "// differentiability witness for "
+     << demangleSymbol(originalFunction->getName()) << "\n";
   // sil_differentiability_witness @original-function-name : $original-sil-type
   PrintOptions qualifiedSILTypeOptions = PrintOptions::printQualifiedSILType();
   OS << "sil_differentiability_witness ";
   printLinkage(OS, linkage, ForDefinition);
-  OS << "@" << originalFunction->getName() << " : "
-     << originalFunction->getLoweredType();
-  // parameters (0, 1, ...)
-  OS << " parameters (";
+  // [parameters 0 1 ...]
+  OS << "[parameters ";
   interleave(parameterIndices->getIndices(),
              [&](unsigned index) { OS << index; },
-             [&] { OS << ", "; });
-  // results (0, 1, ...)
-  OS << ") results (";
+             [&] { OS << " "; });
+  // [results 0 1 ...]
+  OS << "] [results ";
   interleave(resultIndices->getIndices(),
              [&](unsigned index) { OS << index; },
-             [&] { OS << ", "; });
-  OS << ')';
-  // wrt 0, 1, ...
+             [&] { OS << " "; });
+  OS << ']';
+  // [where ...]
   if (derivativeGenericSignature) {
     // NOTE: This needs to be changed if there is no utility for parsing
     // generic signatures. Idea: we could instead print the type of the original
@@ -3096,16 +3096,20 @@ void SILDifferentiabilityWitness::print(
       }
     }
     if (!requirements.empty()) {
-      OS << " where ";
+      OS << " [where ";
       auto SubPrinter = PrintOptions::printSIL();
       interleave(requirements,
                  [&](Requirement req) {
                    req.print(OS, SubPrinter);
                    return;
                  },
                  [&] { OS << ", "; });
+      OS << ']';
     }
   }
+  // original: @original-function-name : $original-sil-type
+  OS << " @" << originalFunction->getName() << " : "
+     << originalFunction->getLoweredType();
   // {
   //   jvp: @jvp-function-name : $jvp-sil-type
   //   vjp: @vjp-function-name : $vjp-sil-type

lib/SIL/SILVerifier.cpp

@@ -5358,14 +5358,16 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const {
     return;
 #endif
   auto origFnType = originalFunction->getLoweredFunctionType();
+  CanGenericSignature derivativeCanGenSig;
+  if (auto *derivativeGenSig = getDerivativeGenericSignature())
+    derivativeCanGenSig = derivativeGenSig->getCanonicalSignature();
   if (jvp) {
     // TODO(TF-893): Change `SILFunctionType::getAutoDiffDerivativeFunctionType`
     // to accept result indices.
     auto expectedJVPType = origFnType->getAutoDiffDerivativeFunctionType(
         getParameterIndices(), /*resultIndex*/ *resultIndices->begin(),
         AutoDiffDerivativeFunctionKind::JVP, M.Types,
-        LookUpConformanceInModule(M.getSwiftModule()),
-        getDerivativeGenericSignature()->getCanonicalSignature());
+        LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig);
     SILVerifier(*jvp).requireSameType(
         SILType::getPrimitiveObjectType(jvp->getLoweredFunctionType()),
         SILType::getPrimitiveObjectType(expectedJVPType),
@@ -5377,8 +5379,7 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const {
     auto expectedVJPType = origFnType->getAutoDiffDerivativeFunctionType(
         getParameterIndices(), /*resultIndex*/ *resultIndices->begin(),
         AutoDiffDerivativeFunctionKind::VJP, M.Types,
-        LookUpConformanceInModule(M.getSwiftModule()),
-        getDerivativeGenericSignature()->getCanonicalSignature());
+        LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig);
     SILVerifier(*vjp).requireSameType(
         SILType::getPrimitiveObjectType(vjp->getLoweredFunctionType()),
         SILType::getPrimitiveObjectType(expectedVJPType),

test/AutoDiff/sil_differentiability_witness_parse.sil

@@ -133,12 +133,13 @@ bb0(%0 : $*T.TangentVector, %1 : $*T.TangentVector):
 // static AdditiveArithmetic<>.zero.getter
 sil [serialized] [always_inline] @$ss18AdditiveArithmeticPss27ExpressibleByIntegerLiteralRzrlE4zeroxvgZ : $@convention(method) <τ_0_0 where τ_0_0 : AdditiveArithmetic, τ_0_0 : ExpressibleByIntegerLiteral> (@thick τ_0_0.Type) -> @out τ_0_0
 
-sil_differentiability_witness hidden @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 parameters (0, 1) results (0) where τ_0_0 : _Differentiable {
+sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {
   jvp: @AD__foo__jvp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, Float) -> @out τ_0_0.TangentVector)
   vjp: @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, Float))
 }
 
-// CHECK-LABEL: sil_differentiability_witness hidden @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 parameters (0, 1) results (0) where τ_0_0 : _Differentiable {
+// CHECK-LABEL: // differentiability witness for foo
+// CHECK: sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {
 // CHECK:   jvp: @AD__foo__jvp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, Float) -> @out τ_0_0.TangentVector)
 // CHECK:   vjp: @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, Float))
 // CHECK: }