[AutoDiff] Diagnose unsupported forward-mode control flow. (#27684)

Diagnose unsupported forward-mode control flow instead of crashing.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -513,6 +513,8 @@ NOTE(autodiff_class_member_not_supported,none,
 NOTE(autodiff_cannot_param_subset_thunk_partially_applied_orig_fn,none,
      "cannot convert a direct method reference to a '@differentiable' "
      "function; use an explicit closure instead", ())
+NOTE(autodiff_jvp_control_flow_not_supported,none,
+     "forward-mode differentiation does not yet support control flow", ())
 NOTE(autodiff_control_flow_not_supported,none,
      "cannot differentiate unsupported control flow", ())
 // TODO(TF-645): Remove when differentiation supports `ref_element_addr`.

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -8030,6 +8030,15 @@ bool ADContext::processDifferentiableAttribute(
     // generation because generated JVP may not match semantics of custom VJP.
     // Instead, create an empty JVP.
     if (RunJVPGeneration && !vjp) {
+      // JVP and differential generation do not currently support functions with
+      // multiple basic blocks.
+      if (original->getBlocks().size() > 1) {
+        emitNondifferentiabilityError(
+            original->getLocation().getSourceLoc(), invoker,
+            diag::autodiff_jvp_control_flow_not_supported);
+        return true;
+      }
+
       JVPEmitter emitter(*this, original, attr, jvp, invoker);
       if (emitter.run())
         return true;

test/AutoDiff/forward_mode_diagnostics.swift

@@ -94,3 +94,18 @@ func nondiff(_ f: @differentiable (Float, @nondiff Float) -> Float) -> Float {
   // expected-error @+1 {{function is not differentiable}}
   return derivative(at: 2, 3) { (x, y) in f(x * x, y) }
 }
+
+//===----------------------------------------------------------------------===//
+// Control flow
+//===----------------------------------------------------------------------===//
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable
+// expected-note @+2 {{when differentiating this function definition}}
+// expected-note @+1 {{forward-mode differentiation does not yet support control flow}}
+func cond(_ x: Float) -> Float {
+  if x > 0 {
+    return x * x
+  }
+  return x + x
+}