[AutoDiff upstream] Add SIL differentiability witness IRGen. (#29704)

dan-zheng · web-flow · commit a17424315983 · 2020-02-07T14:10:34.000-08:00
SIL differentiability witnesses are a new top-level SIL construct mapping an "original" SIL function and derivative configuration to derivative SIL functions. This patch adds `SILDifferentiabilityWitness` IRGen. `SILDifferentiabilityWitness` has a fixed `{ i8*, i8* }` layout: JVP and VJP derivative function pointers. Resolves TF-1146.
diff --git a/include/swift/AST/PrettyStackTrace.h b/include/swift/AST/PrettyStackTrace.h
@@ -202,6 +202,24 @@ class PrettyStackTraceSelector : public llvm::PrettyStackTraceEntry {
   void print(llvm::raw_ostream &OS) const override;
 };
 
+/// PrettyStackTraceDifferentiabilityWitness - Observe that we are processing a
+/// specific differentiability witness.
+class PrettyStackTraceDifferentiabilityWitness
+    : public llvm::PrettyStackTraceEntry {
+  const SILDifferentiabilityWitnessKey Key;
+  const char *Action;
+
+public:
+  PrettyStackTraceDifferentiabilityWitness(
+      const char *action, const SILDifferentiabilityWitnessKey key)
+      : Key(key), Action(action) {}
+  virtual void print(llvm::raw_ostream &OS) const;
+};
+
+void printDifferentiabilityWitnessDescription(
+    llvm::raw_ostream &out, const SILDifferentiabilityWitnessKey key,
+    bool addNewline = true);
+
 } // end namespace swift
 
 #endif
diff --git a/include/swift/IRGen/Linking.h b/include/swift/IRGen/Linking.h
@@ -346,6 +346,9 @@ class LinkEntity {
     /// ProtocolConformance*.
     ProtocolWitnessTableLazyCacheVariable,
 
+    /// A SIL differentiability witness.
+    DifferentiabilityWitness,
+
     // Everything following this is a type kind.
 
     /// A value witness for a type.
@@ -535,6 +538,14 @@ class LinkEntity {
     return getAssociatedConformanceByIndex(conformance->getProtocol(), index);
   }
 
+  void
+  setForDifferentiabilityWitness(Kind kind,
+                                 const SILDifferentiabilityWitness *witness) {
+    Pointer = const_cast<void *>(static_cast<const void *>(witness));
+    SecondaryPointer = nullptr;
+    Data = LINKENTITY_SET_FIELD(Kind, unsigned(kind));
+  }
+
   void setForType(Kind kind, CanType type) {
     assert(isTypeKind(kind));
     Pointer = type.getPointer();
@@ -835,6 +846,14 @@ class LinkEntity {
     return entity;
   }
 
+  static LinkEntity
+  forDifferentiabilityWitness(const SILDifferentiabilityWitness *witness) {
+    LinkEntity entity;
+    entity.setForDifferentiabilityWitness(Kind::DifferentiabilityWitness,
+                                          witness);
+    return entity;
+  }
+
   static LinkEntity forProtocolWitnessTable(const RootProtocolConformance *C) {
     LinkEntity entity;
     entity.setForProtocolConformance(Kind::ProtocolWitnessTable, C);
@@ -1043,6 +1062,11 @@ class LinkEntity {
     return reinterpret_cast<SILGlobalVariable*>(Pointer);
   }
 
+  SILDifferentiabilityWitness *getSILDifferentiabilityWitness() const {
+    assert(getKind() == Kind::DifferentiabilityWitness);
+    return reinterpret_cast<SILDifferentiabilityWitness *>(Pointer);
+  }
+
   const RootProtocolConformance *getRootProtocolConformance() const {
     assert(isRootProtocolConformanceKind(getKind()));
     return cast<RootProtocolConformance>(getProtocolConformance());
diff --git a/lib/AST/AutoDiff.cpp b/lib/AST/AutoDiff.cpp
@@ -17,6 +17,18 @@
 
 using namespace swift;
 
+void AutoDiffConfig::print(llvm::raw_ostream &s) const {
+  s << "(parameters=";
+  parameterIndices->print(s);
+  s << " results=";
+  resultIndices->print(s);
+  if (derivativeGenericSignature) {
+    s << " where=";
+    derivativeGenericSignature->print(s);
+  }
+  s << ')';
+}
+
 // TODO(TF-874): This helper is inefficient and should be removed. Unwrapping at
 // most once (for curried method types) is sufficient.
 static void unwrapCurryLevels(AnyFunctionType *fnTy,
diff --git a/lib/AST/PrettyStackTrace.cpp b/lib/AST/PrettyStackTrace.cpp
@@ -273,3 +273,18 @@ void PrettyStackTraceGenericSignature::print(llvm::raw_ostream &out) const {
 void PrettyStackTraceSelector::print(llvm::raw_ostream &out) const {
   out << "While " << Action << " '" << Selector << "'";
 }
+
+void PrettyStackTraceDifferentiabilityWitness::print(
+    llvm::raw_ostream &out) const {
+  out << "While " << Action << ' ';
+  printDifferentiabilityWitnessDescription(out, Key);
+}
+
+void swift::printDifferentiabilityWitnessDescription(
+    llvm::raw_ostream &out, const SILDifferentiabilityWitnessKey key,
+    bool addNewline) {
+  out << key.first << " ";
+  key.second.print(out);
+  if (addNewline)
+    out << '\n';
+}
diff --git a/lib/IRGen/CMakeLists.txt b/lib/IRGen/CMakeLists.txt
@@ -16,6 +16,7 @@ add_swift_host_library(swiftIRGen STATIC
   GenControl.cpp
   GenCoverage.cpp
   GenDecl.cpp
+  GenDiffWitness.cpp
   GenEnum.cpp
   GenExistential.cpp
   GenFunc.cpp
diff --git a/lib/IRGen/GenDecl.cpp b/lib/IRGen/GenDecl.cpp
@@ -1073,7 +1073,21 @@ void IRGenerator::emitGlobalTopLevel(llvm::StringSet<> *linkerDirectives) {
     CurrentIGMPtr IGM = getGenModule(prop.getDecl()->getInnermostDeclContext());
     IGM->emitSILProperty(&prop);
   }
-  
+
+  // Emit differentiability witnesses.
+  for (auto &dw :
+           PrimaryIGM->getSILModule().getDifferentiabilityWitnessList()) {
+    // Emit into same IRGenModule as the original function.
+    // NOTE(TF-894): Investigate whether `getGenModule(dw.getVJP())` is
+    // significant/desirable; `getGenModule` seems relevant for multi-threaded
+    // compilation. When the differentiation transform canonicalizes all
+    // differentiability witnesses to have JVP/VJP functions, we can assert
+    // that JVP/VJP functions exist and use `getGenModule(dw.getVJP())`.
+    CurrentIGMPtr IGM = getGenModule(dw.getOriginalFunction());
+
+    IGM->emitSILDifferentiabilityWitness(&dw);
+  }
+
   // Emit code coverage mapping data.
   PrimaryIGM->emitCoverageMapping();
 
@@ -4495,6 +4509,13 @@ IRGenModule::getAddrOfWitnessTablePattern(const NormalProtocolConformance *conf,
   return getAddrOfLLVMVariable(entity, definition, DebugTypeInfo());
 }
 
+/// Look up the address of a differentiability witness.
+llvm::Constant *IRGenModule::getAddrOfDifferentiabilityWitness(
+    const SILDifferentiabilityWitness *witness, ConstantInit definition) {
+  auto entity = LinkEntity::forDifferentiabilityWitness(witness);
+  return getAddrOfLLVMVariable(entity, definition, DebugTypeInfo());
+}
+
 llvm::Function *
 IRGenModule::getAddrOfAssociatedTypeWitnessTableAccessFunction(
                                   const NormalProtocolConformance *conformance,
diff --git a/lib/IRGen/GenDiffWitness.cpp b/lib/IRGen/GenDiffWitness.cpp
@@ -0,0 +1,54 @@
+//===--- GenDiffWitness.cpp - IRGen for differentiability witnesses -------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2020 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
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements IR generation for SIL differentiability witnesses.
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/AST/PrettyStackTrace.h"
+#include "swift/SIL/SILDifferentiabilityWitness.h"
+
+#include "ConstantBuilder.h"
+#include "IRGenModule.h"
+
+using namespace swift;
+using namespace irgen;
+
+void IRGenModule::emitSILDifferentiabilityWitness(
+    SILDifferentiabilityWitness *dw) {
+  PrettyStackTraceDifferentiabilityWitness _st(
+      "emitting differentiability witness for", dw->getKey());
+
+  // Don't emit declarations.
+  if (dw->isDeclaration())
+    return;
+
+  // Don't emit `public_external` witnesses.
+  if (dw->getLinkage() == SILLinkage::PublicExternal)
+    return;
+
+  ConstantInitBuilder builder(*this);
+  auto diffWitnessContents = builder.beginStruct();
+
+  assert(dw->getJVP() &&
+         "Differentiability witness definition should have JVP");
+  assert(dw->getVJP() &&
+         "Differentiability witness definition should have VJP");
+
+  diffWitnessContents.addBitCast(
+      getAddrOfSILFunction(dw->getJVP(), NotForDefinition), Int8PtrTy);
+  diffWitnessContents.addBitCast(
+      getAddrOfSILFunction(dw->getVJP(), NotForDefinition), Int8PtrTy);
+
+  getAddrOfDifferentiabilityWitness(
+      dw, diffWitnessContents.finishAndCreateFuture());
+}
diff --git a/lib/IRGen/IRGenModule.cpp b/lib/IRGen/IRGenModule.cpp
@@ -524,6 +524,9 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
 
   DynamicReplacementKeyTy = createStructType(*this, "swift.dyn_repl_key",
                                              {RelativeAddressTy, Int32Ty});
+
+  DifferentiabilityWitnessTy = createStructType(
+      *this, "swift.differentiability_witness", {Int8PtrTy, Int8PtrTy});
 }
 
 IRGenModule::~IRGenModule() {
diff --git a/lib/IRGen/IRGenModule.h b/lib/IRGen/IRGenModule.h
@@ -95,6 +95,7 @@ namespace swift {
   class RootProtocolConformance;
   struct SILDeclRef;
   class SILDefaultWitnessTable;
+  class SILDifferentiabilityWitness;
   class SILGlobalVariable;
   class SILModule;
   class SILProperty;
@@ -672,6 +673,8 @@ class IRGenModule {
       *DynamicReplacementLinkEntryPtrTy; // %link_entry*
   llvm::StructType *DynamicReplacementKeyTy; // { i32, i32}
 
+  llvm::StructType *DifferentiabilityWitnessTy; // { i8*, i8* }
+
   llvm::GlobalVariable *TheTrivialPropertyDescriptor = nullptr;
 
   /// Used to create unique names for class layout types with tail allocated
@@ -1272,6 +1275,7 @@ private:                            \
   void emitSILFunction(SILFunction *f);
   void emitSILWitnessTable(SILWitnessTable *wt);
   void emitSILProperty(SILProperty *prop);
+  void emitSILDifferentiabilityWitness(SILDifferentiabilityWitness *dw);
   void emitSILStaticInitializers();
   llvm::Constant *emitFixedTypeLayout(CanType t, const FixedTypeInfo &ti);
   void emitProtocolConformance(const ConformanceDescription &record);
@@ -1463,6 +1467,10 @@ private:                            \
   llvm::Function *getAddrOfDefaultAssociatedConformanceAccessor(
                                            AssociatedConformance requirement);
 
+  llvm::Constant *
+  getAddrOfDifferentiabilityWitness(const SILDifferentiabilityWitness *witness,
+                                    ConstantInit definition = ConstantInit());
+
   Address getAddrOfObjCISAMask();
 
   /// Retrieve the generic signature for the current generic context, or null if no
diff --git a/lib/IRGen/Linking.cpp b/lib/IRGen/Linking.cpp
@@ -414,6 +414,10 @@ std::string LinkEntity::mangleAsString() const {
   case Kind::ReflectionAssociatedTypeDescriptor:
     return mangler.mangleReflectionAssociatedTypeDescriptor(
                                                     getProtocolConformance());
+  case Kind::DifferentiabilityWitness:
+    return mangler.mangleSILDifferentiabilityWitnessKey(
+        {getSILDifferentiabilityWitness()->getOriginalFunction()->getName(),
+         getSILDifferentiabilityWitness()->getConfig()});
   }
   llvm_unreachable("bad entity kind!");
 }
@@ -659,6 +663,8 @@ SILLinkage LinkEntity::getLinkage(ForDefinition_t forDefinition) const {
   case Kind::ExtensionDescriptor:
   case Kind::AnonymousDescriptor:
     return SILLinkage::Shared;
+  case Kind::DifferentiabilityWitness:
+    return getSILDifferentiabilityWitness()->getLinkage();
   }
   llvm_unreachable("bad link entity kind");
 }
@@ -783,6 +789,9 @@ bool LinkEntity::isAvailableExternally(IRGenModule &IGM) const {
                                      ->getDeclContext()
                                      ->getInnermostTypeContext());
   }
+
+  case Kind::DifferentiabilityWitness:
+    return true;
   
   case Kind::ObjCMetadataUpdateFunction:
   case Kind::ObjCResilientClassStub:
@@ -904,6 +913,8 @@ llvm::Type *LinkEntity::getDefaultDeclarationType(IRGenModule &IGM) const {
       return IGM.ObjCResilientClassStubTy;
     }
     llvm_unreachable("invalid metadata address");
+  case Kind::DifferentiabilityWitness:
+    return IGM.DifferentiabilityWitnessTy;
   default:
     llvm_unreachable("declaration LLVM type not specified");
   }
@@ -951,6 +962,7 @@ Alignment LinkEntity::getAlignment(IRGenModule &IGM) const {
   case Kind::OpaqueTypeDescriptorAccessorKey:
   case Kind::OpaqueTypeDescriptorAccessorVar:
   case Kind::ObjCResilientClassStub:
+  case Kind::DifferentiabilityWitness:
     return IGM.getPointerAlignment();
   case Kind::TypeMetadataDemanglingCacheVariable:
     return Alignment(8);
@@ -1052,6 +1064,7 @@ bool LinkEntity::isWeakImported(ModuleDecl *module) const {
   case Kind::ReflectionBuiltinDescriptor:
   case Kind::ReflectionFieldDescriptor:
   case Kind::CoroutineContinuationPrototype:
+  case Kind::DifferentiabilityWitness:
     return false;
   }
 
@@ -1181,6 +1194,7 @@ const SourceFile *LinkEntity::getSourceFileForEmission() const {
   case Kind::ReflectionBuiltinDescriptor:
   case Kind::ValueWitness:
   case Kind::ValueWitnessTable:
+  case Kind::DifferentiabilityWitness:
     return nullptr;
   }
   
diff --git a/test/AutoDiff/SIL/Serialization/sil_differentiability_witness.sil b/test/AutoDiff/SIL/Serialization/sil_differentiability_witness.sil
@@ -7,11 +7,14 @@
 // RUN: %empty-directory(%t)
 // RUN: %target-sil-opt %s -emit-sib -o %t/tmp.sib -module-name main
 // RUN: %target-sil-opt %t/tmp.sib -o %t/tmp.sil -module-name main
-
 // NOTE(SR-12090): Workaround because import declarations are not preserved in .sib files.
 // RUN: sed -e 's/import Swift$/import Swift; import _Differentiation/' %t/tmp.sil > %t/tmp_fixed.sil
 // RUN: %target-sil-opt %t/tmp_fixed.sil -module-name main -emit-sorted-sil | %FileCheck --check-prefix=ROUNDTRIP %s
 
+// IRGen test.
+
+// RUN: %target-swift-frontend -emit-ir %s | %FileCheck --check-prefix=IRGEN %s
+
 // REQUIRES: differentiable_programming
 // NOTE(SR-12090): `shell` is required only to run `sed` as a SR-12090 workaround.
 // REQUIRES: shell
@@ -49,6 +52,11 @@ sil_differentiability_witness [parameters 0] [results 0] @externalFn1 : $@conven
 // ROUNDTRIP:   vjp: @AD__externalFn1__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
 // ROUNDTRIP: }
 
+// IRGEN-LABEL: @AD__externalFn1_PSRS ={{( protected)?}} global { i8*, i8* } {
+// IRGEN-SAME: @AD__externalFn1__jvp_src_0_wrt_0
+// IRGEN-SAME: @AD__externalFn1__vjp_src_0_wrt_0
+// IRGEN-SAME: }
+
 // Test SIL differentiability witness for bodiless original function, with bodiless jvp/vjp.
 
 sil @externalFn2 : $@convention(thin) (Float) -> Float
@@ -68,6 +76,11 @@ sil_differentiability_witness [parameters 0] [results 0] @externalFn2 : $@conven
 // ROUNDTRIP:   vjp: @AD__externalFn2__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
 // ROUNDTRIP: }
 
+// IRGEN-LABEL: @AD__externalFn2_PSRS ={{( protected)?}} global { i8*, i8* } {
+// IRGEN-SAME: @AD__externalFn2__jvp_src_0_wrt_0
+// IRGEN-SAME: @AD__externalFn2__vjp_src_0_wrt_0
+// IRGEN-SAME: }
+
 // Test SIL differentiability witness declaration.
 
 sil @externalFn3 : $@convention(thin) (Float) -> Float
@@ -77,6 +90,8 @@ sil_differentiability_witness [parameters 0] [results 0] @externalFn3 : $@conven
 // ROUNDTRIP-LABEL: // differentiability witness for externalFn3
 // ROUNDTRIP: sil_differentiability_witness{{( public_external)?}} [parameters 0] [results 0] @externalFn3 : $@convention(thin) (Float) -> Float{{[^{]*$}}
 
+// IRGEN-NOT: @AD__externalFn3{{.*}}={{.*}}{ i8*, i8* }
+
 // Test public non-generic function.
 // SIL differentiability witness:
 // - Has public linkage (implicit).
@@ -108,6 +123,11 @@ sil_differentiability_witness [parameters 0] [results 0] @foo : $@convention(thi
 // ROUNDTRIP:   vjp: @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
 // ROUNDTRIP: }
 
+// IRGEN-LABEL: @AD__foo_PSRS ={{( protected)?}} global { i8*, i8* } {
+// IRGEN-SAME: @AD__foo__jvp_src_0_wrt_0
+// IRGEN-SAME: @AD__foo__vjp_src_0_wrt_0
+// IRGEN-SAME: }
+
 // Test internal generic function.
 // SIL differentiability witness:
 // - Has hidden linkage.
@@ -140,3 +160,8 @@ sil_differentiability_witness hidden [parameters 0 1] [results 0] <τ_0_0 where
 // ROUNDTRIP:   jvp: @AD__generic__jvp_src_0_wrt_0_1 : $@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)
 // ROUNDTRIP:   vjp: @AD__generic__vjp_src_0_wrt_0_1 : $@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))
 // ROUNDTRIP: }
+
+// IRGEN-LABEL: @AD__generic_PSSRS16_Differentiation14DifferentiableRzl = hidden global { i8*, i8* } {
+// IRGEN-SAME: @AD__generic__jvp_src_0_wrt_0_1
+// IRGEN-SAME: @AD__generic__vjp_src_0_wrt_0_1
+// IRGEN-SAME: }

Original file line number	Diff line number	Diff line change
`@@ -524,6 +524,9 @@ IRGenModule::IRGenModule(IRGenerator &irgen,`
`524`	`524`
`525`	`525`	`DynamicReplacementKeyTy = createStructType(*this, "swift.dyn_repl_key",`
`526`	`526`	`{RelativeAddressTy, Int32Ty});`
	`527`	`+`
	`528`	`+ DifferentiabilityWitnessTy = createStructType(`
	`529`	`+ *this, "swift.differentiability_witness", {Int8PtrTy, Int8PtrTy});`
`527`	`530`	`}`
`528`	`531`
`529`	`532`	`IRGenModule::~IRGenModule() {`