[autodiff] When asserts are enabled, verify all autodiff compiler generated functions.

This ensures that any invalid SIL generated by these cloners is caught
immediately at the source when asserts are enabled improving productivity.

Author: gottesmm

include/swift/SILOptimizer/Differentiation/JVPCloner.h

@@ -47,6 +47,8 @@ class JVPCloner final {
   /// Performs JVP generation on the empty JVP function. Returns true if any
   /// error occurs.
   bool run();
+
+  SILFunction &getJVP() const;
 };
 
 } // end namespace autodiff

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -31,6 +31,9 @@
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/DenseMap.h"
 
+using namespace swift;
+using namespace autodiff;
+
 namespace swift {
 namespace autodiff {
 
@@ -380,6 +383,8 @@ class JVPCloner::Implementation final
   /// Run JVP generation. Returns true on error.
   bool run();
 
+  SILFunction &getJVP() const { return *jvp; }
+
   void postProcess(SILInstruction *orig, SILInstruction *cloned) {
     if (errorOccurred)
       return;
@@ -1727,7 +1732,16 @@ bool JVPCloner::Implementation::run() {
   return errorOccurred;
 }
 
-bool JVPCloner::run() { return impl.run(); }
-
 } // end namespace autodiff
 } // end namespace swift
+
+bool JVPCloner::run() {
+  bool foundError = impl.run();
+#ifndef NDEBUG
+  if (!foundError)
+    getJVP().verify();
+#endif
+  return foundError;
+}
+
+SILFunction &JVPCloner::getJVP() const { return impl.getJVP(); }

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -138,7 +138,6 @@ class PullbackCloner::Implementation final
   SILModule &getModule() const { return getContext().getModule(); }
   ASTContext &getASTContext() const { return getPullback().getASTContext(); }
   SILFunction &getOriginal() const { return vjpCloner.getOriginal(); }
-  SILFunction &getPullback() const { return vjpCloner.getPullback(); }
   SILDifferentiabilityWitness *getWitness() const {
     return vjpCloner.getWitness();
   }
@@ -782,6 +781,10 @@ class PullbackCloner::Implementation final
   /// parameters.
   void emitZeroDerivativesForNonvariedResult(SILValue origNonvariedResult);
 
+  /// Public helper so that our users can get the underlying newly created
+  /// function.
+  SILFunction &getPullback() const { return vjpCloner.getPullback(); }
+
   using TrampolineBlockSet = SmallPtrSet<SILBasicBlock *, 4>;
 
   /// Determines the pullback successor block for a given original block and one
@@ -1740,7 +1743,14 @@ PullbackCloner::~PullbackCloner() { delete &impl; }
 // Entry point
 //--------------------------------------------------------------------------//
 
-bool PullbackCloner::run() { return impl.run(); }
+bool PullbackCloner::run() {
+  bool foundError = impl.run();
+#ifndef NDEBUG
+  if (!foundError)
+    impl.getPullback().verify();
+#endif
+  return foundError;
+}
 
 bool PullbackCloner::Implementation::run() {
   PrettyStackTraceSILFunction trace("generating pullback for", &getOriginal());

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -1019,7 +1019,14 @@ bool VJPCloner::Implementation::run() {
   return errorOccurred;
 }
 
-bool VJPCloner::run() { return impl.run(); }
+bool VJPCloner::run() {
+  bool foundError = impl.run();
+#ifndef NDEBUG
+  if (!foundError)
+    getVJP().verify();
+#endif
+  return foundError;
+}
 
 } // end namespace autodiff
 } // end namespace swift