Merge pull request #34887 from gottesmm/pr-a2f43c89854b9be02db9525f4baea04ff3bca6b8

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

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