[AutoDiff] Disable differentiable_function_extract explicit type as… (#35239)

`differentiability_function_extract` instruction has an optional explicit
extractee type. This is currently used by TypeSubstCloner and the
LoadableByAddress transform to rewrite `differentiability_function_extract`
instructions while preserving `@differentiable` function type invariants.

There is an assertion that `differentiability_function_extract` instructions do
not have explicit extractee types outside of canonical/lowered SIL. However,
this does not handle the SIL deserialization case above: when a function
containing a `differentiable_function_extract` instruction with an explicit type
is deserialized into a raw SIL module (which happens when optimizations are
enabled).

Removing the assertion unblocks this encountered use case.

A more robust longer-term solution may be to change SIL `@differentiable`
function types to explicitly store component original/JVP/VJP function types.

Also fix `differentiable_function_extract` extractee type serialization.

Resolves SR-14004.

Author: dan-zheng

include/swift/SIL/SILInstruction.h

@@ -8920,7 +8920,12 @@ class DifferentiableFunctionExtractInst
                    SILModule &module);
 
 public:
-  /// Note: explicit extractee type may be specified only in lowered SIL.
+  /// Note: explicit extractee type is used to avoid inconsistent typing in:
+  /// - Canonical SIL, due to generic specialization.
+  /// - Lowered SIL, due to LoadableByAddress.
+  /// - Raw SIL, due to deserialization of canonical/lowered SIL functions.
+  /// See `TypeSubstCloner::visitDifferentiableFunctionExtractInst` for an
+  /// explanation of how explicit extractee type is used.
   explicit DifferentiableFunctionExtractInst(
       SILModule &module, SILDebugLocation debugLoc,
       NormalDifferentiableFunctionTypeComponent extractee, SILValue function,

lib/SIL/IR/SILInstructions.cpp

@@ -731,18 +731,6 @@ DifferentiableFunctionExtractInst::DifferentiableFunctionExtractInst(
                                : getExtracteeType(function, extractee, module),
                            function.getOwnershipKind()),
       Extractee(extractee), HasExplicitExtracteeType(extracteeType.hasValue()) {
-#ifndef NDEBUG
-  if (extracteeType.hasValue()) {
-    // Note: explicit extractee type is used to avoid inconsistent typing in:
-    // - Canonical SIL, due to generic specialization.
-    // - Lowered SIL, due to LoadableByAddress.
-    // See `TypeSubstCloner::visitDifferentiableFunctionExtractInst` for an
-    // explanation of how explicit extractee type is used.
-    assert((module.getStage() == SILStage::Canonical ||
-            module.getStage() == SILStage::Lowered) &&
-           "Explicit type is valid only in canonical or lowered SIL");
-  }
-#endif
 }
 
 SILType LinearFunctionExtractInst::

lib/Serialization/DeserializeSIL.cpp

@@ -1186,7 +1186,7 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn,
   case SIL_INST_DIFFERENTIABLE_FUNCTION_EXTRACT:
     SILInstDifferentiableFunctionExtractLayout::readRecord(
         scratch, TyID, TyCategory, ValID, /*extractee*/ Attr,
-        /*hasExplicitExtracteeType*/ Attr2);
+        /*hasExplicitExtracteeType*/ Attr2, /*explicitExtracteeType*/ TyID2);
     RawOpCode = (unsigned)SILInstructionKind::DifferentiableFunctionExtractInst;
     break;
   case SIL_INST_LINEAR_FUNCTION_EXTRACT:
@@ -2747,8 +2747,11 @@ bool SILDeserializer::readSILInstruction(SILFunction *Fn,
     auto val = getLocalValue(ValID, silTy);
     NormalDifferentiableFunctionTypeComponent extractee(Attr);
     Optional<SILType> explicitExtracteeType = None;
-    if (Attr2)
-      explicitExtracteeType = silTy;
+    if (Attr2) {
+      auto extracteeASTType = MF->getType(TyID2);
+      explicitExtracteeType =
+          getSILType(extracteeASTType, SILValueCategory::Object, Fn);
+    }
     ResultInst = Builder.createDifferentiableFunctionExtract(
         Loc, extractee, val, explicitExtracteeType);
     break;

lib/Serialization/ModuleFormat.h

@@ -56,7 +56,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 589; // cache prespecialization decls.
+const uint16_t SWIFTMODULE_VERSION_MINOR = 590; // differentiable_function_extract explicit extractee type
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///

lib/Serialization/SILFormat.h

@@ -478,8 +478,9 @@ namespace sil_block {
     TypeIDField,
     SILTypeCategoryField,
     ValueIDField,
-    BCFixed<2>, // extractee
-    BCFixed<1>  // has explicit extractee type?
+    BCFixed<2>,  // extractee
+    BCFixed<1>,  // has explicit extractee type?
+    TypeIDField  // explicit extractee type
   >;
 
   using SILInstLinearFunctionExtractLayout = BCRecordLayout<

lib/Serialization/SerializeSIL.cpp

@@ -2294,11 +2294,13 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
     auto operandType = dfei->getOperand()->getType();
     auto operandTypeRef = S.addTypeRef(operandType.getASTType());
     auto rawExtractee = (unsigned)dfei->getExtractee();
+    auto extracteeTypeRef = S.addTypeRef(dfei->getType().getASTType());
     SILInstDifferentiableFunctionExtractLayout::emitRecord(
         Out, ScratchRecord,
         SILAbbrCodes[SILInstDifferentiableFunctionExtractLayout::Code],
         operandTypeRef, (unsigned)operandType.getCategory(), operandRef,
-        rawExtractee, (unsigned)dfei->hasExplicitExtracteeType());
+        rawExtractee, (unsigned)dfei->hasExplicitExtracteeType(),
+        extracteeTypeRef);
     break;
   }
   case SILInstructionKind::LinearFunctionExtractInst: {

test/AutoDiff/compiler_crashers_fixed/sr14004-cross-module-differentiation-differentiable-function-extract-inlining.swift

@@ -0,0 +1,34 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift-dylib(%t/%target-library-name(Library)) -emit-module -emit-module-path %t/Library.swiftmodule -module-name Library -DLIBRARY %s
+// RUN: %target-build-swift -I %t -O -emit-module %s
+
+// SR-14004: Assertion failure due to function with `differentiable_function_extract`
+// with explicit extractee type being deserialized into a raw SIL module.
+
+#if LIBRARY
+
+import _Differentiation
+
+public struct Struct<Scalar>: Differentiable {}
+
+@differentiable
+public func foo<Scalar>(_ x: Struct<Scalar>) -> Struct<Scalar> { x }
+
+@inlinable
+@differentiable
+public func bar<Scalar>(_ x: Struct<Scalar>) -> Struct<Scalar> {
+  foo(x)
+}
+
+#else
+
+import _Differentiation
+import Library
+
+public func foo(
+  body: @differentiable (Struct<Float>) -> Struct<Float> = bar
+) {
+  fatalError()
+}
+
+#endif