Finish parsing/printing/serialization.

Author: dan-zheng

include/swift/AST/DiagnosticsParse.def

@@ -686,6 +686,20 @@ ERROR(sil_witness_assoc_conf_not_found,none,
 ERROR(sil_witness_protocol_conformance_not_found,none,
       "sil protocol conformance not found", ())
 
+// [differentiable ...] (sil-decl attr)
+ERROR(sil_diff_witness_expected_keyword,PointsToFirstBadToken,
+      "expected '%0' in differentiability witness", (StringRef))
+ERROR(sil_diff_witness_expected_parameter_list,PointsToFirstBadToken,
+      "expected an comma-separated list of parameter indices, e.g. (0, 1)", ())
+ERROR(sil_diff_witness_expected_rsquare,PointsToFirstBadToken,
+      "expected ']' to end 'differentiable' attribute", ())
+ERROR(sil_diff_witness_expected_parameter_index,PointsToFirstBadToken,
+      "expected the index of a parameter to differentiate w.r.t.", ())
+ERROR(sil_diff_witness_expected_source_index,PointsToFirstBadToken,
+      "expected the index of a result to differentiate from", ())
+
+// SIL differentiability witnesses
+
 // SIL Coverage Map
 ERROR(sil_coverage_func_not_found, none,
       "sil function not found %0", (Identifier))

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -42,8 +42,10 @@ class SILDifferentiabilityWitness
       public SILAllocated<SILDifferentiabilityWitness>
 {
 private:
-  /// The module which contains the SIL differentiability witness.
+  /// The module which contains the differentiability witness.
   SILModule &module;
+  /// The linkage of the differentiability witness.
+  SILLinkage linkage;
   /// The original function.
   SILFunction *originalFunction;
   /// The parameter indices.
@@ -60,19 +62,31 @@ class SILDifferentiabilityWitness
   /// devirtualization from another module.
   bool serialized;
 
-  SILDifferentiabilityWitness(SILModule &module, SILFunction *originalFunction,
+  SILDifferentiabilityWitness(SILModule &module, SILLinkage linkage,
+                              SILFunction *originalFunction,
                               AutoDiffIndexSubset *parameterIndices,
                               AutoDiffIndexSubset *resultIndices,
                               GenericSignature *derivativeGenSig,
                               SILFunction *jvp, SILFunction *vjp,
                               bool isSerialized)
-    : module(module), originalFunction(originalFunction),
+    : module(module), linkage(linkage), originalFunction(originalFunction),
       parameterIndices(parameterIndices), resultIndices(resultIndices),
       derivativeGenericSignature(derivativeGenSig), jvp(jvp), vjp(vjp),
       serialized(isSerialized) {}
 
 public:
+  /// The key type, used for uniquing `SILDifferentiabilityWitness` in
+  /// `SILModule`, original function, parameter indices, result indices, and
+  /// derivative generic signature.
+  using Key = std::tuple<const SILFunction *, AutoDiffIndexSubset *,
+                         AutoDiffIndexSubset *, GenericSignature *>;
+  Key getKey() {
+    return std::make_tuple(originalFunction, parameterIndices, resultIndices,
+                           derivativeGenericSignature);
+  }
+
   SILModule &getModule() const { return module; }
+  SILLinkage getLinkage() const { return linkage; }
   SILFunction *getOriginalFunction() const { return originalFunction; }
   AutoDiffIndexSubset *getParameterIndices() const {
     return parameterIndices;
@@ -88,7 +102,7 @@ class SILDifferentiabilityWitness
   bool isSerialized() const { return serialized; }
 
   static SILDifferentiabilityWitness *create(
-      SILModule &module, SILFunction *originalFunction,
+      SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
       AutoDiffIndexSubset *parameterIndices, AutoDiffIndexSubset *resultIndices,
       GenericSignature *derivativeGenSig, SILFunction *jvp, SILFunction *vjp,
       bool isSerialized);

include/swift/SIL/SILModule.h

@@ -145,6 +145,9 @@ class SILModule {
   friend SILProperty;
   friend SILUndef;
   friend SILWitnessTable;
+  // SWIFT_ENABLE_TENSORFLOW
+  friend SILDifferentiabilityWitness;
+  // SWIFT_ENABLE_TENSORFLOW END
   friend Lowering::SILGenModule;
   friend Lowering::TypeConverter;
   class SerializationCallback;
@@ -202,10 +205,9 @@ class SILModule {
 
   // SWIFT_ENABLE_TENSORFLOW
   /// Lookup table for SIL differentiability witnesses from original functions.
-  /// Indexed by original function, parameter indices, result indices, and
-  /// derivative generic signature.
-  llvm::DenseMap<std::tuple<const SILFunction *, AutoDiffIndexSubset *,
-                            AutoDiffIndexSubset *, GenericSignature *>,
+  /// Indexed by key type: original function, parameter indices, result indices,
+  /// and derivative generic signature.
+  llvm::DenseMap<SILDifferentiabilityWitness::Key,
                  SILDifferentiabilityWitness *>
   DifferentiabilityWitnessMap;
 

include/swift/Serialization/SerializedSILLoader.h

@@ -99,6 +99,11 @@ class SerializedSILLoader {
   /// Deserialize all Properties in all SILModules.
   void getAllProperties();
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Deserialize all DifferentiabilityWitnesses in all SILModules.
+  void getAllDifferentiabilityWitnesses();
+  // SWIFT_ENABLE_TENSORFLOW END
+
   SerializedSILLoader(const SerializedSILLoader &) = delete;
   SerializedSILLoader(SerializedSILLoader &&) = delete;
   SerializedSILLoader &operator=(const SerializedSILLoader &) = delete;

lib/Parse/ParseDecl.cpp

@@ -94,6 +94,9 @@ bool Parser::parseTopLevel() {
     CASE_SIL(sil_global, SILGlobal)
     CASE_SIL(sil_witness_table, SILWitnessTable)
     CASE_SIL(sil_default_witness_table, SILDefaultWitnessTable)
+    // SWIFT_ENABLE_TENSORFLOW
+    CASE_SIL(sil_differentiability_witness, SILDifferentiabilityWitness)
+    // SWIFT_ENABLE_TENSORFLOW END
     CASE_SIL(sil_coverage_map, SILCoverageMap)
     CASE_SIL(sil_property, SILProperty)
     CASE_SIL(sil_scope, SILScope)

lib/ParseSIL/ParseSIL.cpp

@@ -6745,21 +6745,225 @@ bool SILParserTUState::parseSILDefaultWitnessTable(Parser &P) {
   return false;
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+// TODO: Dedupe with `SILParser::convertRequirements` upstream.
+// Consider defining this as `Parser::convertRequirements`.
+static void convertRequirements(Parser &P, SILFunction *F,
+                                ArrayRef<RequirementRepr> From,
+                                SmallVectorImpl<Requirement> &To) {
+  if (From.empty()) {
+    To.clear();
+    return;
+  }
+
+  auto *GenericEnv = F->getLoweredFunctionType()
+                         ->getGenericSignature()
+                         ->getGenericEnvironment();
+  assert(GenericEnv);
+  (void)GenericEnv;
+
+  IdentTypeReprLookup PerformLookup(P);
+  // Use parser lexical scopes to resolve references
+  // to the generic parameters.
+  auto ResolveToInterfaceType = [&](TypeLoc Ty) -> Type {
+    Ty.getTypeRepr()->walk(PerformLookup);
+    swift::performTypeLocChecking(P.Context, Ty, /*isSILMode*/ true,
+                                  /*isSILType*/ true, GenericEnv, &P.SF);
+    assert(Ty.getType());
+    return Ty.getType()->mapTypeOutOfContext();
+  };
+
+  for (auto &Req : From) {
+    if (Req.getKind() == RequirementReprKind::SameType) {
+      auto FirstType = ResolveToInterfaceType(Req.getFirstTypeLoc());
+      auto SecondType = ResolveToInterfaceType(Req.getSecondTypeLoc());
+      Requirement ConvertedRequirement(RequirementKind::SameType, FirstType,
+                                       SecondType);
+      To.push_back(ConvertedRequirement);
+      continue;
+    }
+
+    if (Req.getKind() == RequirementReprKind::TypeConstraint) {
+      auto Subject = ResolveToInterfaceType(Req.getSubjectLoc());
+      auto Constraint = ResolveToInterfaceType(Req.getConstraintLoc());
+      Requirement ConvertedRequirement(RequirementKind::Conformance, Subject,
+                                       Constraint);
+      To.push_back(ConvertedRequirement);
+      continue;
+    }
+
+    if (Req.getKind() == RequirementReprKind::LayoutConstraint) {
+      auto Subject = ResolveToInterfaceType(Req.getSubjectLoc());
+      Requirement ConvertedRequirement(RequirementKind::Layout, Subject,
+                                       Req.getLayoutConstraint());
+      To.push_back(ConvertedRequirement);
+      continue;
+    }
+    llvm_unreachable("Unsupported requirement kind");
+  }
+}
+
 /// decl-sil-differentiability-witness ::=
 ///   'sil_differentiability_witness'
-///   sil-function-name
-///   'wrt' autodiff-index-subset
-///   'sources' autodiff-index-subset
-///   ('derivative_generic_signature' generic-signature)?
-///   '{' ('jvp' sil-function-name)? ('vjp' sil-function-name)? '}'
+///   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]+)*
+///   '(' [0-9]+ (',', [0-9]+)* ')'
 bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   P.consumeToken(tok::kw_sil_differentiability_witness);
-  // TODO(TF-867): Implement parsing. Test round-tripping with printing.
+  SILParser State(P);
+
+  // Parse the linkage.
+  Optional<SILLinkage> linkage;
+  if (parseSILLinkage(linkage, P))
+    return true;
+  if (!linkage)
+    linkage = SILLinkage::PublicExternal;
+
+  Scope S(&P, ScopeKind::TopLevel);
+  Scope Body(&P, ScopeKind::FunctionBody);
+
+  auto parseFunctionNameAndType = [&](SILFunction *&fn) -> bool {
+    Identifier name;
+    SILType ty;
+    SourceLoc fnNameLoc = P.Tok.getLoc();
+    // We need to turn on InSILBody to parse the function reference.
+    Lexer::SILBodyRAII tmp(*P.L);
+    GenericEnvironment *ignoredEnv;
+    if ((State.parseGlobalName(name)) ||
+        P.parseToken(tok::colon, diag::expected_sil_colon_value_ref) ||
+        State.parseSILType(ty, ignoredEnv, /*IsFuncDecl*/ true))
+      return true;
+
+    // The function doesn't exist yet. Create a zombie forward declaration.
+    auto fnType = ty.getAs<SILFunctionType>();
+    if (!fnType || !ty.isObject()) {
+      P.diagnose(fnNameLoc, diag::expected_sil_function_type);
+      return true;
+    }
+    fn = State.getGlobalNameForReference(name, fnType, fnNameLoc, true);
+    State.TUState.PotentialZombieFns.insert(fn);
+    return false;
+  };
+
+  SourceLoc lastLoc = P.getEndOfPreviousLoc();
+
+  SILFunction *originalFn;
+  if (parseFunctionNameAndType(originalFn))
+    return true;
+
+  auto parseAutoDiffIndexSubset =
+      [&](StringRef label, AutoDiffIndexSubset *& paramIndexSubset) -> bool {
+    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.
+    auto parseParam = [&]() -> bool {
+      unsigned index;
+      // TODO: Reject non-ascending parameter index lists.
+      if (P.parseUnsignedInteger(index, lastLoc,
+              diag::sil_diff_witness_expected_parameter_list))
+        return true;
+      paramIndices.push_back(index);
+      return false;
+    };
+    // Parse first.
+    if (parseParam())
+      return true;
+    // Parse rest.
+    while (P.consumeIf(tok::comma))
+      if (parseParam())
+        return true;
+    if (P.parseToken(tok::r_paren, diag::sil_diff_witness_expected_keyword,
+                     "("))
+      return true;
+    auto maxIndexRef =
+        std::max_element(paramIndices.begin(), paramIndices.end());
+    paramIndexSubset = AutoDiffIndexSubset::get(
+        P.Context, maxIndexRef ? *maxIndexRef + 1 : 0, paramIndices);
+    return false;
+  };
+  AutoDiffIndexSubset *parameterIndices = nullptr;
+  AutoDiffIndexSubset *resultIndices = nullptr;
+  if (parseAutoDiffIndexSubset("parameters", parameterIndices))
+    return true;
+  if (parseAutoDiffIndexSubset("results", resultIndices))
+    return true;
+
+  GenericSignature *derivativeGenSig = nullptr;
+  // Parse a trailing 'where' clause if any.
+  if (P.Tok.is(tok::kw_where)) {
+    SourceLoc whereLoc;
+    SmallVector<RequirementRepr, 4> requirementReprs;
+    bool firstTypeInComplete;
+    P.parseGenericWhereClause(whereLoc, requirementReprs, firstTypeInComplete,
+                              /*AllowLayoutConstraints*/ false);
+    auto *whereClause = TrailingWhereClause::create(
+        originalFn->getModule().getASTContext(), whereLoc, requirementReprs);
+    SmallVector<Requirement, 4> requirements;
+    convertRequirements(P, originalFn, whereClause->getRequirements(),
+                        requirements);
+    assert(requirements.size() == requirementReprs.size());
+    derivativeGenSig = evaluateOrDefault(
+        P.Context.evaluator,
+        AbstractGenericSignatureRequest{
+            originalFn->getLoweredFunctionType()->getGenericSignature(),
+            /*addedGenericParams=*/{},
+            std::move(requirements)},
+            nullptr);
+  }
+
+  SILFunction *jvp = nullptr;
+  SILFunction *vjp = nullptr;
+  if (P.Tok.is(tok::l_brace)) {
+    SourceLoc LBraceLoc = P.Tok.getLoc();
+    P.consumeToken(tok::l_brace);
+
+    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "jvp") {
+      P.consumeToken(tok::identifier);
+      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_keyword,
+                       ":"))
+        return true;
+      Scope Body(&P, ScopeKind::FunctionBody);
+      if (parseFunctionNameAndType(jvp))
+        return true;
+    }
+
+    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "vjp") {
+      P.consumeToken(tok::identifier);
+      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_keyword,
+                       ":"))
+        return true;
+      Scope Body(&P, ScopeKind::FunctionBody);
+      if (parseFunctionNameAndType(vjp))
+        return true;
+    }
+
+    if (P.parseMatchingToken(tok::r_brace, lastLoc, diag::expected_sil_rbrace,
+                             LBraceLoc))
+      return true;
+  }
+
+  // TODO: Parse `isSerialized` flag.
+  bool isSerialized = false;
+  SILDifferentiabilityWitness::create(
+      M, *linkage, originalFn, parameterIndices, resultIndices,
+      derivativeGenSig, jvp, vjp, isSerialized);
   return false;
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 llvm::Optional<llvm::coverage::Counter> SILParser::parseSILCoverageExpr(
     llvm::coverage::CounterExpressionBuilder &Builder) {

lib/SIL/SILDifferentiabilityWitness.cpp

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -18,15 +18,19 @@
 using namespace swift;
 
 SILDifferentiabilityWitness *SILDifferentiabilityWitness::create(
-    SILModule &module, SILFunction *originalFunction,
+    SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
     AutoDiffIndexSubset *parameterIndices, AutoDiffIndexSubset *resultIndices,
     GenericSignature *derivativeGenSig, SILFunction *jvp, SILFunction *vjp,
     bool isSerialized) {
   void *buf = module.allocate(sizeof(SILDifferentiabilityWitness),
                               alignof(SILDifferentiabilityWitness));
-  SILDifferentiabilityWitness *dw = ::new (buf)
-      SILDifferentiabilityWitness(module, originalFunction, parameterIndices,
-                                  resultIndices, derivativeGenSig, jvp, vjp,
-                                  isSerialized);
-  return dw;
+  auto *diffWitness = ::new (buf) SILDifferentiabilityWitness(
+      module, linkage, originalFunction, parameterIndices, resultIndices,
+      derivativeGenSig, jvp, vjp, isSerialized);
+  // Register the differentiability witness in the module.
+  assert(!module.DifferentiabilityWitnessMap.count(diffWitness->getKey()) &&
+         "Cannot create duplicate differentiability witness in a module");
+  module.DifferentiabilityWitnessMap[diffWitness->getKey()] = diffWitness;
+  module.getDifferentiabilityWitnessList().push_back(diffWitness);
+  return diffWitness;
 }