[AutoDiff][IRGen] Differentiation of optional initializers

[philipturner]

Describe the solution you'd like
The Swift Stdlib currently lacks support for differentiating conversion of types to optionals. #57228 was a feature request about converting Double to Double?. I did some digging into the proposed feature, and encountered a compiler crasher. This regards the operator ?? (or perhaps it was !), which I attempted to replicate using an instance method test instead of an explicit operator.

To incorporate this feature into the Standard Library, one would need to modify the code below. There could be a way to work around the crash. In this instance, the compiler crashed while compiling philipturner/differentiation on the v5.6 release toolchain. One may have different results using the v5.7 release toolchain offered in Xcode 14 beta, or when using a development toolchain. Someone should also try adding this code to the Swift Stdlib's source files and ensuring apple/swift successfully compiles.

After finding a suitable version of the reproducer, locate the test file tests/AutoDiff/stdlib/optional.swift. One should add a new code section that shows Optional.init being differentiated.

Current reproducer:

extension Optional where Wrapped: Differentiable {
  @usableFromInline
  @derivative(of: init(_:))
  static func _vjpInit(_ value: Wrapped) -> (
    value: Wrapped?, pullback: (TangentVector) -> Wrapped.TangentVector
  ) {
    return (Self(value), {
      $0.value! // ?? Wrapped.TangentVector.zero
    })
  }

  @usableFromInline
  @derivative(of: init(_:))
  static func _jvpInit(_ value: Wrapped) -> (
    value: Wrapped?, differential: (Wrapped.TangentVector) -> TangentVector
  ) {
    return (Self(value), {
      TangentVector($0) // ?? Wrapped.TangentVector.zero
    })
  }
  
  @inlinable
  static func test(optional: Wrapped?, defaultValue: @autoclosure () throws -> Wrapped) rethrows -> Wrapped {
    fatalError()
  }
  
  @usableFromInline
  @derivative(of: test, wrt: `optional`)
  static func _vjpForceUnwrap(optional: Wrapped?, defaultValue: @autoclosure () throws -> Wrapped) rethrows -> (
    value: Wrapped, pullback: (Wrapped.TangentVector) -> TangentVector
  ) {
    fatalError()
  }
}

Crash:

(base) philipturner@m1-max-mbp differentiation % swift build
warning: Usage of /Users/philipturner/Library/org.swift.swiftpm/collections.json has been deprecated. Please delete it and use the new /Users/philipturner/Library/org.swift.swiftpm/configuration/collections.json instead.
Building for debugging...
Please submit a bug report (https://swift.org/contributing/#reporting-bugs) and include the project and the crash backtrace.
Stack dump:
0.	Program arguments: /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/swift-frontend -frontend -c /Users/philipturner/Desktop/differentiation/Sources/Differentiation/AnyDifferentiable.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/ArrayDifferentiation.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/DifferentialOperators.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/Differentiation.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/DifferentiationUtilities.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/FloatingPointDifferentiation.swift -primary-file /Users/philipturner/Desktop/differentiation/Sources/Differentiation/OptionalDifferentiation.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/SIMDDifferentiation.swift /Users/philipturner/Desktop/differentiation/Sources/Differentiation/TgmathDerivatives.swift -emit-dependencies-path /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug/Differentiation.build/OptionalDifferentiation.d -emit-reference-dependencies-path /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug/Differentiation.build/OptionalDifferentiation.swiftdeps -disable-objc-attr-requires-foundation-module -target arm64-apple-macosx10.10 -Xllvm -aarch64-use-tbi -enable-objc-interop -sdk /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.3.sdk -I /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug -I /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/usr/lib -F /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -color-diagnostics -enable-testing -g -module-cache-path /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug/ModuleCache -parse-stdlib -swift-version 5 -Onone -D SWIFT_PACKAGE -D DEBUG -new-driver-path /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/swift-driver -resource-dir /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib/swift -enable-anonymous-context-mangled-names -module-name Differentiation -target-sdk-version 12.3 -parse-as-library -o /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug/Differentiation.build/OptionalDifferentiation.swift.o -index-store-path /Users/philipturner/Desktop/differentiation/.build/arm64-apple-macosx/debug/index/store -index-system-modules
1.	Apple Swift version 5.6 (swiftlang-5.6.0.323.62 clang-1316.0.20.8)
2.	Compiling with the current language version
3.	While evaluating request IRGenRequest(IR Generation for file "/Users/philipturner/Desktop/differentiation/Sources/Differentiation/OptionalDifferentiation.swift")
4.	While emitting IR SIL function "@$sSq15Differentiation01_A014DifferentiableRzlE4test8optionalxxSg_tFZAbCRzlTJrSUpSr".
 for '_vjpForceUnwrap(optional:)' (at /Users/philipturner/Desktop/differentiation/Sources/Differentiation/OptionalDifferentiation.swift:96:3)
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  swift-frontend           0x0000000108c06f88 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
1  swift-frontend           0x0000000108c05f9c llvm::sys::RunSignalHandlers() + 112
2  swift-frontend           0x0000000108c07618 SignalHandler(int) + 344
3  libsystem_platform.dylib 0x000000019cd194c4 _sigtramp + 56
4  swift-frontend           0x0000000104978ca0 swift::irgen::IRGenModule::emitSILFunction(swift::SILFunction*) + 1884
5  swift-frontend           0x0000000104978ca0 swift::irgen::IRGenModule::emitSILFunction(swift::SILFunction*) + 1884
6  swift-frontend           0x000000010484f010 swift::irgen::IRGenerator::emitGlobalTopLevel(std::__1::vector<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> >, std::__1::allocator<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > > const&) + 1560
7  swift-frontend           0x000000010494b060 swift::IRGenRequest::evaluate(swift::Evaluator&, swift::IRGenDescriptor) const + 7140
8  swift-frontend           0x0000000104978278 swift::SimpleRequest<swift::IRGenRequest, swift::GeneratedModule (swift::IRGenDescriptor), (swift::RequestFlags)9>::evaluateRequest(swift::IRGenRequest const&, swift::Evaluator&) + 188
9  swift-frontend           0x0000000104955c40 llvm::Expected<swift::IRGenRequest::OutputType> swift::Evaluator::getResultUncached<swift::IRGenRequest>(swift::IRGenRequest const&) + 736
10 swift-frontend           0x000000010494fc30 swift::performIRGeneration(swift::FileUnit*, swift::IRGenOptions const&, swift::TBDGenOptions const&, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::StringRef, swift::PrimarySpecificPaths const&, llvm::StringRef, llvm::GlobalVariable**) + 236
11 swift-frontend           0x00000001044b516c performCompileStepsPostSILGen(swift::CompilerInstance&, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::PointerUnion<swift::ModuleDecl*, swift::SourceFile*>, swift::PrimarySpecificPaths const&, int&, swift::FrontendObserver*) + 2524
12 swift-frontend           0x00000001044a8840 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 5344
13 swift-frontend           0x0000000104470130 swift::mainEntry(int, char const**) + 808
14 dyld                     0x000000010b48d088 start + 516

Related: SR-13440, SR-13627, SR-14881

Worst case scenario

There is no way to differentiate Optional.init in AutoDiff's current state. One may have to modify the C++ code and permit differentiating throwing functions. Alternatively, one may hunt down the bug causing the compiler crash above. If the bug's source if found, there can be a PR to apple/swift detailing:

• (a) The C++ code modifications that fix the bug.
• (b) A compiler crasher test involving optional differentiation.
• (c) The addition of Optional.init to the Swift Stdlib.
• (d) Additions to the test/AutoDiff/stdlib suite.

Other details

This is not directly in the interest of apple/swift, but I would like to validate s4tf/s4tf against this addition to the Swift Stdilb. I recently faced #59467, which caused a lot of long-term difficulties. If this causes a compiler crash while building S4TF or philipturner/differentiation, I would request further investigation before merging.