[AutoDiff] Use ASTMangler for generated AD associated functions. (#26389)

dan-zheng · web-flow · commit cf28c000c13d · 2019-07-29T12:16:30.000-07:00
Use ASTMangler for generated AD associated functions and linear maps: JVPs, VJPs, differentials, pullbacks. This centralizes ad-hoc mangling logic scattered throughout the codebase. Add `AutoDiffLinearMapKind` enum. Todos: - TF-680: Mangle `@differentiable` attribute requirements. - TF-685: Use ASTMangler for AD-related thunks. - TF-686: Make Demangler/Remangler work with AD-generated functions.
diff --git a/include/swift/AST/ASTMangler.h b/include/swift/AST/ASTMangler.h
@@ -151,7 +151,25 @@ class ASTMangler : public Mangler {
                                              Type FromType, Type ToType,
                                              Type SelfType,
                                              ModuleDecl *Module);
-  
+
+  // SWIFT_ENABLE_TENSORFLOW
+  // Mangle the autodiff associated function (JVP/VJP) with the given:
+  // - Mangled original function name.
+  // - Associated function kind.
+  // - Parameter/result indices.
+  std::string mangleAutoDiffAssociatedFunctionHelper(
+      StringRef name, AutoDiffAssociatedFunctionKind kind,
+      const SILAutoDiffIndices &indices);
+
+  // SWIFT_ENABLE_TENSORFLOW
+  // Mangle the autodiff linear map (differential/pullback) with the given:
+  // - Mangled original function name.
+  // - Linear map kind.
+  // - Parameter/result indices.
+  std::string mangleAutoDiffLinearMapHelper(
+      StringRef name, AutoDiffLinearMapKind kind,
+      const SILAutoDiffIndices &indices);
+
   std::string mangleKeyPathGetterThunkHelper(const AbstractStorageDecl *property,
                                              GenericSignature *signature,
                                              CanType baseType,
diff --git a/include/swift/AST/AutoDiff.h b/include/swift/AST/AutoDiff.h
@@ -546,6 +546,20 @@ struct AutoDiffAssociatedFunctionKind {
   operator innerty() const { return rawValue; }
 };
 
+/// The kind of an linear map.
+struct AutoDiffLinearMapKind {
+  enum innerty : uint8_t {
+     // The differential function.
+     Differential = 0,
+     // The pullback function.
+     Pullback = 1
+  } rawValue;
+
+  AutoDiffLinearMapKind() = default;
+  AutoDiffLinearMapKind(innerty rawValue) : rawValue(rawValue) {}
+  operator innerty() const { return rawValue; }
+};
+
 /// In conjunction with the original function decl, identifies an associated
 /// autodiff function.
 ///
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -369,6 +369,52 @@ std::string ASTMangler::mangleReabstractionThunkHelper(
   return finalize();
 }
 
+std::string ASTMangler::mangleAutoDiffAssociatedFunctionHelper(
+    StringRef name, AutoDiffAssociatedFunctionKind kind,
+    const SILAutoDiffIndices &indices) {
+  // TODO(TF-20): Make the mangling scheme robust.
+  // TODO(TF-680): Mangle `@differentiable` atttribute requirements as well.
+  beginManglingWithoutPrefix();
+
+  Buffer << "AD__" << name << '_';
+  switch (kind) {
+  case AutoDiffAssociatedFunctionKind::JVP:
+    Buffer << "_jvp_";
+    break;
+  case AutoDiffAssociatedFunctionKind::VJP:
+    Buffer << "_vjp_";
+    break;
+  }
+  Buffer << indices.mangle();
+
+  auto result = Storage.str().str();
+  Storage.clear();
+  return result;
+}
+
+std::string ASTMangler::mangleAutoDiffLinearMapHelper(
+    StringRef name, AutoDiffLinearMapKind kind,
+    const SILAutoDiffIndices &indices) {
+  // TODO(TF-20): Make the mangling scheme robust.
+  // TODO(TF-680): Mangle `@differentiable` atttribute requirements as well.
+  beginManglingWithoutPrefix();
+
+  Buffer << "AD__" << name << '_';
+  switch (kind) {
+  case AutoDiffLinearMapKind::Differential:
+    Buffer << "_differential_";
+    break;
+  case AutoDiffLinearMapKind::Pullback:
+    Buffer << "_pullback_";
+    break;
+  }
+  Buffer << indices.mangle();
+
+  auto result = Storage.str().str();
+  Storage.clear();
+  return result;
+}
+
 std::string ASTMangler::mangleTypeForDebugger(Type Ty, const DeclContext *DC) {
   PrettyStackTraceType prettyStackTrace(Ty->getASTContext(),
                                         "mangling type for debugger", Ty);
diff --git a/lib/SIL/SILDeclRef.cpp b/lib/SIL/SILDeclRef.cpp
@@ -677,6 +677,9 @@ static void mangleClangDecl(raw_ostream &buffer,
 }
 
 std::string SILDeclRef::mangle(ManglingKind MKind) const {
+  using namespace Mangle;
+  ASTMangler mangler;
+
   // SWIFT_ENABLE_TENSORFLOW
   if (autoDiffAssociatedFunctionIdentifier) {
     std::string originalMangled = asAutoDiffOriginalFunction().mangle(MKind);
@@ -686,22 +689,11 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
         autoDiffAssociatedFunctionIdentifier->getParameterIndices()->getLowered(
             functionTy->getASTContext(), functionTy);
     SILAutoDiffIndices indices(/*source*/ 0, silParameterIndices);
-    std::string mangledKind;
-    switch (autoDiffAssociatedFunctionIdentifier->getKind()) {
-    case AutoDiffAssociatedFunctionKind::JVP:
-      mangledKind = "jvp";
-      break;
-    case AutoDiffAssociatedFunctionKind::VJP:
-      mangledKind = "vjp";
-      break;
-    }
-    return "AD__" + originalMangled + "__" + mangledKind + "_" +
-           indices.mangle();
+    auto assocFnKind = autoDiffAssociatedFunctionIdentifier->getKind();
+    return mangler.mangleAutoDiffAssociatedFunctionHelper(
+        originalMangled, assocFnKind, indices);
   }
 
-  using namespace Mangle;
-  ASTMangler mangler;
-
   // As a special case, Clang functions and globals don't get mangled at all.
   if (hasDecl()) {
     if (auto clangDecl = getDecl()->getClangDecl()) {
diff --git a/lib/SIL/SILFunctionBuilder.cpp b/lib/SIL/SILFunctionBuilder.cpp
@@ -13,6 +13,8 @@
 #include "swift/SIL/SILFunctionBuilder.h"
 #include "swift/AST/Availability.h"
 #include "swift/AST/Decl.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/AST/ASTMangler.h"
 using namespace swift;
 
 SILFunction *SILFunctionBuilder::getOrCreateFunction(
@@ -104,12 +106,19 @@ void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
           indices.parameters->getNumIndices() > 1;
       if (isSelfReorderedMethod) {
         auto &ctx = F->getASTContext();
-        if (A->getJVPFunction())
+        if (A->getJVPFunction()) {
+          Mangle::ASTMangler mangler;
           jvpName = ctx.getIdentifier(
-              "AD__" + constant.mangle() + "__jvp_" + indices.mangle()).str();
+              mangler.mangleAutoDiffAssociatedFunctionHelper(
+                  constant.mangle(), AutoDiffAssociatedFunctionKind::JVP,
+                  indices)).str();
+        }
         if (A->getVJPFunction()) {
+          Mangle::ASTMangler mangler;
           vjpName = ctx.getIdentifier(
-              "AD__" + constant.mangle() + "__vjp_" + indices.mangle()).str();
+              mangler.mangleAutoDiffAssociatedFunctionHelper(
+                  constant.mangle(), AutoDiffAssociatedFunctionKind::VJP,
+                  indices)).str();
         }
       } else {
         if (auto *jvpFn = A->getJVPFunction())
diff --git a/lib/SILGen/SILGenPoly.cpp b/lib/SILGen/SILGenPoly.cpp
@@ -3508,18 +3508,10 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionReorderingThunk(
     IsSerialized_t isSerialized) {
   auto assocFnType = assocFn->getLoweredFunctionType();
 
-  std::string name;
-  switch (assocFnKind) {
-  case AutoDiffAssociatedFunctionKind::JVP:
-    name = "jvp";
-    break;
-  case AutoDiffAssociatedFunctionKind::VJP:
-    name = "vjp";
-    break;
-  }
-  name = getASTContext().getIdentifier(
-      "AD__" + original->getName().str() + "__" + name + "_" +
-      indices.mangle()).str();
+  Mangle::ASTMangler mangler;
+  auto name = getASTContext().getIdentifier(
+      mangler.mangleAutoDiffAssociatedFunctionHelper(
+          original->getName(), assocFnKind, indices)).str();
 
   Lowering::GenericContextScope genericContextScope(
       Types, assocFnType->getGenericSignature());
diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp
@@ -2057,7 +2057,7 @@ emitAssociatedFunctionReference(
     auto substMap = getSubstitutionMap(original);
     // Attempt to look up a `[differentiable]` attribute that minimally
     // satisfies the specified indices.
-    // TODO(TF-482): Change `lookupMinimalDifferentiableAttr` to additionally
+    // TODO(TF-482): Change `lookUpMinimalDifferentiableAttr` to additionally
     // check whether `[differentiable]` attribute generic requirements are
     // satisfied.
     auto *minimalAttr =
@@ -2112,7 +2112,7 @@ emitAssociatedFunctionReference(
     }
     assert(minimalAttr);
     // TODO(TF-482): Move generic requirement checking logic to
-    // `lookupMinimalDifferentiableAttr`.
+    // `lookUpMinimalDifferentiableAttr`.
     if (!checkRequirementsSatisfied(
             minimalAttr->getRequirements(),
             substMap, originalFn, context.getModule().getSwiftModule())) {
@@ -2959,10 +2959,11 @@ class VJPEmitter final
               ->getCanonicalType(), origParam.getConvention()));
     }
 
-    auto pbName = original->getASTContext()
-                       .getIdentifier("AD__" + original->getName().str() +
-                                      "__pullback_" + indices.mangle())
-                       .str();
+    Mangle::ASTMangler mangler;
+    auto pbName = original->getASTContext().getIdentifier(
+        mangler.mangleAutoDiffLinearMapHelper(
+            original->getName(), AutoDiffLinearMapKind::Pullback,
+            indices)).str();
     auto pbGenericSig = getAssociatedFunctionGenericSignature(attr, original);
     auto *pbGenericEnv = pbGenericSig
         ? pbGenericSig->createGenericEnvironment()
@@ -3460,7 +3461,7 @@ class VJPEmitter final
   void visitAutoDiffFunctionInst(AutoDiffFunctionInst *adfi) {
     // Clone `autodiff_function` from original to VJP, then add the cloned
     // instruction to the `autodiff_function` worklist.
-    SILClonerWithScopes::visitAutoDiffFunctionInst(adfi);
+    TypeSubstCloner::visitAutoDiffFunctionInst(adfi);
     auto *newADFI = cast<AutoDiffFunctionInst>(getOpValue(adfi));
     context.getAutoDiffFunctionInsts().push_back(newADFI);
   }
@@ -3550,11 +3551,11 @@ class JVPEmitter final
               ->getAutoDiffAssociatedTangentSpace(lookupConformance)
               ->getCanonicalType(), origParam.getConvention()));
     }
-
-    auto diffName = original->getASTContext()
-        .getIdentifier("AD__" + original->getName().str() + "__differential_" +
-                       indices.mangle())
-        .str();
+    Mangle::ASTMangler mangler;
+    auto diffName = original->getASTContext().getIdentifier(
+        mangler.mangleAutoDiffLinearMapHelper(
+            original->getName(), AutoDiffLinearMapKind::Differential,
+            indices)).str();
     auto diffGenericSig = getAssociatedFunctionGenericSignature(attr, original);
     auto *diffGenericEnv = diffGenericSig
         ? diffGenericSig->createGenericEnvironment()
@@ -5918,10 +5919,11 @@ static SILFunction *createEmptyVJP(
   auto indices = attr->getIndices();
 
   // === Create an empty VJP. ===
-  auto vjpName = original->getASTContext()
-                     .getIdentifier("AD__" + original->getName().str() +
-                                    "__vjp_" + indices.mangle())
-                     .str();
+  Mangle::ASTMangler mangler;
+  auto vjpName = original->getASTContext().getIdentifier(
+      mangler.mangleAutoDiffAssociatedFunctionHelper(
+          original->getName(), AutoDiffAssociatedFunctionKind::VJP, indices))
+              .str();
   auto vjpGenericSig = getAssociatedFunctionGenericSignature(attr, original);
 
   // RAII that pushes the original function's generic signature to
@@ -5968,10 +5970,11 @@ static SILFunction *createEmptyJVP(
   auto indices = attr->getIndices();
 
   // === Create an empty JVP. ===
-  auto jvpName = original->getASTContext()
-      .getIdentifier("AD__" + original->getName().str() +
-                     "__jvp_" + indices.mangle())
-      .str();
+  Mangle::ASTMangler mangler;
+  auto jvpName = original->getASTContext().getIdentifier(
+      mangler.mangleAutoDiffAssociatedFunctionHelper(
+          original->getName(), AutoDiffAssociatedFunctionKind::JVP, indices))
+              .str();
   auto jvpGenericSig = getAssociatedFunctionGenericSignature(attr, original);
 
   // RAII that pushes the original function's generic signature to
@@ -6017,10 +6020,11 @@ bool ADContext::processDifferentiableAttribute(
   if (attr->hasJVP()) {
     jvpName = attr->getJVPName();
   } else if (original->isExternalDeclaration()) {
-    jvpName = original->getASTContext()
-                  .getIdentifier("AD__" + original->getName().str() +
-                                 "__jvp_" + attr->getIndices().mangle())
-                  .str();
+    Mangle::ASTMangler mangler;
+    jvpName = original->getASTContext().getIdentifier(
+        mangler.mangleAutoDiffAssociatedFunctionHelper(
+            original->getName(), AutoDiffAssociatedFunctionKind::JVP,
+            attr->getIndices())).str();
   }
   if (!jvpName.empty()) {
     jvp = module.lookUpFunction(jvpName);
@@ -6044,10 +6048,11 @@ bool ADContext::processDifferentiableAttribute(
   if (attr->hasVJP()) {
     vjpName = attr->getVJPName();
   } else if (original->isExternalDeclaration()) {
-    vjpName = original->getASTContext()
-                  .getIdentifier("AD__" + original->getName().str() +
-				                         "__vjp_" + attr->getIndices().mangle())
-                  .str();
+    Mangle::ASTMangler mangler;
+    vjpName = original->getASTContext().getIdentifier(
+        mangler.mangleAutoDiffAssociatedFunctionHelper(
+            original->getName(), AutoDiffAssociatedFunctionKind::VJP,
+            attr->getIndices())).str();
   }
   if (!vjpName.empty()) {
     vjp = module.lookUpFunction(vjpName);
@@ -6142,7 +6147,7 @@ ADContext::getOrCreateSubsetParametersThunkForLinearMap(
       /*withoutActuallyEscaping*/ true,
       DifferentiationThunkKind::Reabstraction);
 
-  // TODO: Use more principled mangling.
+  // TODO(TF-685): Use more principled mangling for thunks.
   std::string thunkName;
   switch (kind) {
     case AutoDiffAssociatedFunctionKind::JVP:
@@ -6408,7 +6413,7 @@ ADContext::getOrCreateSubsetParametersThunkForAssociatedFunction(
         ->getAbstractFunctionDecl()->getNameStr();
   }
   assert(!origName.empty() && "Original function name could not be resolved");
-  // TODO: Use more principled mangling.
+  // TODO(TF-685): Use more principled mangling for thunks.
   std::string thunkName;
   switch (kind) {
     case AutoDiffAssociatedFunctionKind::JVP:
@@ -6531,6 +6536,7 @@ SILValue ADContext::promoteToDifferentiableFunction(
     if (auto *thunkRef = dyn_cast<FunctionRefInst>(ai->getCallee())) {
       SILAutoDiffIndices desiredIndices(resultIndex, parameterIndices);
       auto *thunk = thunkRef->getReferencedFunctionOrNull();
+      // TODO(TF-685): Use more principled mangling for thunks.
       auto newThunkName = "AD__" + thunk->getName().str() +
           "__differentiable_curry_thunk_" + desiredIndices.mangle();
 
diff --git a/test/AutoDiff/autodiff_indirect_diagnostics.swift b/test/AutoDiff/autodiff_indirect_diagnostics.swift
@@ -35,16 +35,6 @@ func weirdWrapper<T : Differentiable>(_ x: T) -> T {
 }
 _ = gradient(at: Float(1), in: { x in weirdWrapper(x) })
 
-/*
-// FIXME(TF-482): This currently crashes during differentiation transform.
-// because `T` is not constrained to `Differentiable` in generated
-// `[differentiable]` attribute.
-@differentiable
-func directMissingConformance<T>(_ x: T) -> T {
-  return x
-}
-*/
-
 @differentiable
 func direct<T : Differentiable>(_ x: T) -> T {
   return x
diff --git a/test/AutoDiff/differentiable_attr_type_checking.swift b/test/AutoDiff/differentiable_attr_type_checking.swift
@@ -613,6 +613,12 @@ func invalidRequirementLayout<Scalar>(x: Scalar) -> Scalar {
   return x
 }
 
+// expected-error @+1 {{no differentiation parameters could be inferred; must differentiate with respect to at least one parameter conforming to 'Differentiable'}}
+@differentiable
+func missingConformance<T>(_ x: T) -> T {
+  return x
+}
+
 protocol ProtocolRequirements : Differentiable {
   // expected-note @+2 {{protocol requires initializer 'init(x:y:)' with type '(x: Float, y: Float)'}}
   @differentiable
