[AutoDiff] NFC: IRGen gardening. (#28238)

Move code from lib/IRGen/GenProto.cpp to dedicated file
lib/IRGen/GenDiffWitness.cpp.

Author: dan-zheng

lib/IRGen/CMakeLists.txt

@@ -26,6 +26,8 @@ add_swift_host_library(swiftIRGen STATIC
   GenDecl.cpp
   # SWIFT_ENABLE_TENSORFLOW
   GenDiffFunc.cpp
+  GenDiffWitness.cpp
+  # SWIFT_ENABLE_TENSORFLOW END
   GenEnum.cpp
   GenExistential.cpp
   GenFunc.cpp

lib/IRGen/GenDecl.cpp

@@ -1068,7 +1068,7 @@ void IRGenerator::emitGlobalTopLevel() {
   // SWIFT_ENABLE_TENSORFLOW
   // Emit differentiability witnesses.
   for (auto &dw :
-       PrimaryIGM->getSILModule().getDifferentiabilityWitnessList()) {
+           PrimaryIGM->getSILModule().getDifferentiabilityWitnessList()) {
     if (dw.isDeclaration())
       continue;
 

lib/IRGen/GenDiffFunc.cpp

@@ -9,7 +9,13 @@
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
+//
 // SWIFT_ENABLE_TENSORFLOW
+//
+// This file implements IR generation for `@differentiable` function types in
+// Swift.
+//
+//===----------------------------------------------------------------------===//
 
 #include "swift/AST/Decl.h"
 #include "swift/AST/Pattern.h"
@@ -32,7 +38,6 @@
 using namespace swift;
 using namespace irgen;
 
-
 //----------------------------------------------------------------------------//
 // `@differentiable` (non-linear) function type info
 //----------------------------------------------------------------------------//

lib/IRGen/GenDiffWitness.cpp

@@ -0,0 +1,58 @@
+//===--- GenDiffWitness.cpp - IRGen for differentiability witnesses -------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 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
+//
+//===----------------------------------------------------------------------===//
+//
+// SWIFT_ENABLE_TENSORFLOW
+//
+// 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;
+
+  ConstantInitBuilder builder(*this);
+  auto diffWitnessContents = builder.beginStruct();
+
+  // TODO(TF-894): When the differentiation transform canonicalizes all
+  // differentiability witnesses to have JVP/VJP functions, remove the nullptr
+  // cases and assert that JVP/VJP functions exist.
+  if (dw->getJVP()) {
+    diffWitnessContents.addBitCast(
+        getAddrOfSILFunction(dw->getJVP(), NotForDefinition), Int8PtrTy);
+  } else {
+    diffWitnessContents.addNullPointer(Int8PtrTy);
+  }
+  if (dw->getVJP()) {
+    diffWitnessContents.addBitCast(
+        getAddrOfSILFunction(dw->getVJP(), NotForDefinition), Int8PtrTy);
+  } else {
+    diffWitnessContents.addNullPointer(Int8PtrTy);
+  }
+
+  getAddrOfDifferentiabilityWitness(
+      dw, diffWitnessContents.finishAndCreateFuture());
+}

lib/IRGen/GenProto.cpp

@@ -2166,40 +2166,6 @@ void IRGenModule::emitSILWitnessTable(SILWitnessTable *wt) {
                                        RequireMetadata);
 }
 
-// SWIFT_ENABLE_TENSORFLOW
-void IRGenModule::emitSILDifferentiabilityWitness(
-    SILDifferentiabilityWitness *dw) {
-  PrettyStackTraceDifferentiabilityWitness _st(
-      "emitting differentiability witness for", dw->getKey());
-
-  // Don't emit declarations.
-  if (dw->isDeclaration())
-    return;
-
-  ConstantInitBuilder builder(*this);
-  auto diffWitnessContents = builder.beginStruct();
-
-  // TODO(marcrasi): When the differentiation pass generates JVP/VJP for
-  // witnesses, remove the nullptr case and add assertions that the JVP/VJP
-  // exist.
-  if (dw->getJVP()) {
-    diffWitnessContents.addBitCast(
-        getAddrOfSILFunction(dw->getJVP(), NotForDefinition), Int8PtrTy);
-  } else {
-    diffWitnessContents.addNullPointer(Int8PtrTy);
-  }
-  if (dw->getVJP()) {
-    diffWitnessContents.addBitCast(
-        getAddrOfSILFunction(dw->getVJP(), NotForDefinition), Int8PtrTy);
-  } else {
-    diffWitnessContents.addNullPointer(Int8PtrTy);
-  }
-
-  getAddrOfDifferentiabilityWitness(
-      dw, diffWitnessContents.finishAndCreateFuture());
-}
-// SWIFT_ENABLE_TENSORFLOW_END
-
 /// True if a function's signature in LLVM carries polymorphic parameters.
 /// Generic functions and protocol witnesses carry polymorphic parameters.
 bool irgen::hasPolymorphicParameters(CanSILFunctionType ty) {