[AutoDiff upstream] Add forward-mode differentiation. (#30878)

JVP functions are forward-mode derivative functions. They take original
arguments and return original results and a differential function. Differential
functions take derivatives wrt arguments and return derivatives wrt results.

`JVPEmitter` is a cloner that emits JVP and differential functions at the same
time. In JVP functions, function applications are replaced with JVP function
applications. In differential functions, function applications are replaced
with differential function applications.

In JVP functions, each basic block takes a differential struct containing callee
differentials. These structs are consumed by differential functions.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -512,6 +512,8 @@ NOTE(autodiff_cannot_param_subset_thunk_partially_applied_orig_fn,none,
      "function; use an explicit closure instead", ())
 NOTE(autodiff_cannot_differentiate_through_multiple_results,none,
      "cannot differentiate through multiple results", ())
+NOTE(autodiff_cannot_differentiate_through_inout_arguments,none,
+     "cannot differentiate through 'inout' arguments", ())
 // TODO(TF-1149): Remove this diagnostic.
 NOTE(autodiff_loadable_value_addressonly_tangent_unsupported,none,
      "cannot yet differentiate value whose type %0 has a compile-time known "

include/swift/SILOptimizer/Utils/Differentiation/JVPEmitter.h

@@ -449,6 +449,9 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_enable_experimental_additive_arithmetic_derivation))
     Opts.EnableExperimentalAdditiveArithmeticDerivedConformances = true;
 
+  Opts.EnableExperimentalForwardModeDifferentiation |=
+      Args.hasArg(OPT_enable_experimental_forward_mode_differentiation);
+
   Opts.DebuggerSupport |= Args.hasArg(OPT_debugger_support);
   if (Opts.DebuggerSupport)
     Opts.EnableDollarIdentifiers = true;

lib/Frontend/CompilerInvocation.cpp

@@ -39,6 +39,7 @@
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/Differentiation/ADContext.h"
+#include "swift/SILOptimizer/Utils/Differentiation/JVPEmitter.h"
 #include "swift/SILOptimizer/Utils/Differentiation/Thunk.h"
 #include "swift/SILOptimizer/Utils/Differentiation/VJPEmitter.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
@@ -898,9 +899,10 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
          diagnoseUnsupportedControlFlow(context, original, invoker)))
       return true;
 
-    witness->setJVP(
-        createEmptyJVP(context, original, witness, serializeFunctions));
-    context.recordGeneratedFunction(witness->getJVP());
+    // Create empty JVP.
+    auto *jvp = createEmptyJVP(context, original, witness, serializeFunctions);
+    witness->setJVP(jvp);
+    context.recordGeneratedFunction(jvp);
 
     // For now, only do JVP generation if the flag is enabled and if custom VJP
     // does not exist. If custom VJP exists but custom JVP does not, skip JVP
@@ -917,18 +919,18 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
             diag::autodiff_jvp_control_flow_not_supported);
         return true;
       }
-      // TODO(TF-1211): Upstream and use `JVPEmitter`. Fatal error with a nice
-      // message for now.
-      auto *jvp = witness->getJVP();
-      emitFatalError(context, jvp, "_fatalErrorJVPNotGenerated");
+      // Emit JVP function.
+      JVPEmitter emitter(context, original, witness, jvp, invoker);
+      if (emitter.run())
+        return true;
     } else {
       // If JVP generation is disabled or a user-defined custom VJP function
       // exists, fatal error with a nice message.
-      emitFatalError(context, witness->getJVP(),
+      emitFatalError(context, jvp,
                      "_fatalErrorForwardModeDifferentiationDisabled");
       LLVM_DEBUG(getADDebugStream()
                  << "Generated empty JVP for " << original->getName() << ":\n"
-                 << *witness->getJVP());
+                 << *jvp);
     }
   }
 
@@ -945,6 +947,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     auto *vjp = createEmptyVJP(context, original, witness, serializeFunctions);
     witness->setVJP(vjp);
     context.recordGeneratedFunction(vjp);
+    // Emit VJP function.
     VJPEmitter emitter(context, original, witness, vjp, invoker);
     return emitter.run();
   }

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -2,6 +2,7 @@ silopt_register_sources(
   ADContext.cpp
   Common.cpp
   DifferentiationInvoker.cpp
+  JVPEmitter.cpp
   LinearMapInfo.cpp
   PullbackEmitter.cpp
   Thunk.cpp