[AutoDiff] Diagnose non-differentiable original function arguments/results. (#26407)

Diagnose non-differentiable original function arguments/results in
`VJPEmitter::visitApplyInst`. Perform the check again if original function
is specialized via partial application.

Resolves TF-687.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3303,27 +3303,36 @@ class VJPEmitter final
           /*differentiationOrder*/ 1, functionSource);
     }
 
-    // Check and diagnose non-differentiable arguments.
-    for (unsigned paramIndex : range(originalFnTy->getNumParameters())) {
-      if (indices.isWrtParameter(paramIndex) &&
-              !originalFnTy->getParameters()[paramIndex]
-              .getSILStorageType()
-              .isDifferentiable(getModule())) {
-        context.emitNondifferentiabilityError(
-            original, invoker, diag::autodiff_nondifferentiable_argument);
-        errorOccurred = true;
-        return;
-      }
-    }
-    // Check and diagnose non-differentiable results.
-    if (!originalFnTy->getResults()[indices.source]
-            .getSILStorageType()
-            .isDifferentiable(getModule())) {
-      context.emitNondifferentiabilityError(
-          original, invoker, diag::autodiff_nondifferentiable_result);
-      errorOccurred = true;
+    // Check and diagnose non-differentiable original function type.
+    auto diagnoseNondifferentiableOriginalFunctionType =
+        [&](CanSILFunctionType origFnTy) {
+          // Check and diagnose non-differentiable arguments.
+          for (unsigned paramIndex : range(originalFnTy->getNumParameters())) {
+            if (indices.isWrtParameter(paramIndex) &&
+                    !originalFnTy->getParameters()[paramIndex]
+                    .getSILStorageType()
+                    .isDifferentiable(getModule())) {
+              context.emitNondifferentiabilityError(
+                  ai->getArgumentsWithoutIndirectResults()[paramIndex], invoker,
+                  diag::autodiff_nondifferentiable_argument);
+              errorOccurred = true;
+              return true;
+            }
+          }
+          // Check and diagnose non-differentiable results.
+          if (!originalFnTy->getResults()[indices.source]
+                  .getSILStorageType()
+                  .isDifferentiable(getModule())) {
+            context.emitNondifferentiabilityError(
+                original, invoker, diag::autodiff_nondifferentiable_result);
+            errorOccurred = true;
+            return true;
+          }
+          return false;
+        };
+    if (diagnoseNondifferentiableOriginalFunctionType(originalFnTy))
       return;
-    }
+
     // If VJP has not yet been found, emit an `autodiff_function` instruction
     // on the remapped original function operand and `autodiff_function_extract`
     // the VJP. The actual JVP/VJP functions will be populated in the
@@ -3354,6 +3363,10 @@ class VJPEmitter final
             ai->getLoc(), original, substMap, {},
             ParameterConvention::Direct_Guaranteed);
         original = vjpPartialApply;
+        originalFnTy = original->getType().castTo<SILFunctionType>();
+        // Diagnose if new original function type is non-differentiable.
+        if (diagnoseNondifferentiableOriginalFunctionType(originalFnTy))
+          return;
       }
 
       auto *autoDiffFuncInst = context.createAutoDiffFunction(
@@ -3363,6 +3376,8 @@ class VJPEmitter final
 
       // Record the `autodiff_function` instruction.
       context.getAutoDiffFunctionInsts().push_back(autoDiffFuncInst);
+      // TODO(TF-689): Make `autodiff_function` store result indices and remove
+      // `ADContext::resultIndices`.
       context.getResultIndices()[autoDiffFuncInst] =
           activeResultIndices.front();
 

test/AutoDiff/autodiff_diagnostics.swift

@@ -252,6 +252,10 @@ let no_return: @differentiable (Float) -> Float = { x in
 // expected-note @+1 {{missing return for differentiation}}
 }
 
+//===----------------------------------------------------------------------===//
+// Non-differentiable arguments and results
+//===----------------------------------------------------------------------===//
+
 // expected-error @+1 {{function is not differentiable}}
 @differentiable
 // expected-note @+1 {{when differentiating this function definition}}

test/AutoDiff/autodiff_indirect_diagnostics.swift

@@ -70,6 +70,23 @@ struct TF8Struct<Scalar> : TF8Proto where Scalar : FloatingPoint & Differentiabl
 
 _ = gradient(at: 1.0, in: { x in x.squareRoot() })
 
+//===----------------------------------------------------------------------===//
+// Non-differentiable arguments and results
+//===----------------------------------------------------------------------===//
+
+struct TF_687<T> : Differentiable {
+  @noDerivative var indirectDummy: T
+  var base: Float
+
+  init(_ base: Float, dummy: T) {
+    self.base = base
+    self.indirectDummy = dummy
+  }
+}
+// expected-error @+2 {{function is not differentiable}}
+// expected-note @+1 {{cannot differentiate through a non-differentiable argument; do you want to use 'withoutDerivative(at:)'?}}
+let _: @differentiable (Float) -> TF_687<Any> = { x in TF_687<Any>(x, dummy: x) }
+
 //===----------------------------------------------------------------------===//
 // Add `Differentiable` conformance for generic wrt parameters
 //===----------------------------------------------------------------------===//
@@ -87,4 +104,4 @@ extension TF_691: Differentiable where Scalar: Differentiable {}
 
 func identity<T>(_ x: TF_691<T>) -> TF_691<T> { x }
 let _: @differentiable (Float) -> TF_691<Float> = { x in identity(TF_691(x)) }
-let _: @differentiable (Float) -> TF_691<Float> = { x in id(TF_691(x)) }
+let _: @differentiable (Float) -> TF_691<Float> = { x in id(TF_691(x)) }