[AutoDiff][SIL] Compiler crasher in S4TF

[philipturner]

Describe the bug
There is a compiler crasher preventing the default synthesis of callAsFunction in the Layer protocol.

Steps To Reproduce
I don't have the bandwidth to narrow down this compiler crasher, but here's the gist. In tensorflow/swift-apis#1184, the annotation functionality of S4TF was deactivated due to a compiler bug that was never narrowed down. Every layer having a forward function had it renamed to callAsFunction, preventing model annotations from happening.

I tried enabling model annotations, and got multiple disparate compiler crashers. One had the exact same error output as SR-13455 and affected the same section of S4TF code, but I don't know whether it has the same cause. Another was flaky, sometimes requiring that I clean the build folder before it appears again. This flaky crash is what I'm reporting here.

Relevant source code (Sources/TensorFlow/Layer.swift, circa line 166):

public protocol Layer: Module where Input: Differentiable {
  /// Returns the output obtained from applying the layer to the given input.
  ///
  /// - Parameter input: The input to the layer.
  /// - Returns: The output.
  @differentiable(reverse)
  func callAsFunction(_ input: Input) -> Output
  
  @differentiable(reverse)
  func forward(_ input: Input) -> Output
}

extension Layer {
  // Workaround for SR-13455: autodiff undefined symbol linker error.
  @differentiable(reverse, wrt: self)
  @differentiable(reverse)
  public func forward(_ input: Input) -> Output {
    return callAsFunction(input)
  }
}

// Workaround for compiler crasher.
//extension Layer where Input: DifferentiableTensorProtocol, Output: DifferentiableTensorProtocol {
//  @differentiable(reverse)
//  public func callAsFunction(_ input: Input) -> Output {
//    let activation = forward(input)
//    return annotated(activation)
//  }
//}

Relevant source code (Sources/TensorFlow/Layers/Dense.swift, circa line 72):

@frozen
public struct Dense<Scalar: TensorFlowFloatingPoint>: Layer {
  ...
  @differentiable(reverse)
  public func forward(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
    if batched {
      let hidden = matmul(input.expandingShape(at: 1), weight).squeezingShape(at: 1)
      return activation(useBias ? hidden + bias : hidden)
    }
    return activation(useBias ? (matmul(input, weight) + bias) : matmul(input, weight))
  }
  
  // Workaround for compiler crasher.
  @differentiable(reverse)
  public func callAsFunction(_ input: Input) -> Output {
    let activation = forward(input)
    return annotated(activation)
  }
}

Screenshots

The code above is modified so that the crash does not happen. The modifications are highlighted by the comment "Workaround for compiler crasher". Below is the crash that has been suppressed.

Crash

Building for debugging...
Assertion failed: (fn->getLoweredFunctionType() == type), function getOrCreateFunction, file SILFunctionBuilder.cpp, line 32.
Please submit a bug report (https://swift.org/contributing/#reporting-bugs) and include the project and the crash backtrace.
Stack dump:
0.	Program arguments: /Library/Developer/Toolchains/swift-5.7-DEVELOPMENT-SNAPSHOT-2022-06-26-a.xctoolchain/usr/bin/swift-frontend -frontend -c /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/BackwardsCompatibility.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Bindings/EagerExecution.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Bindings/RawOpsAugmented.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Bindings/RawOpsDispatching.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Bindings/RawOpsGenerated.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Bindings/TFTensorOperation.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Context.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/ArrayOps.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/BroadcastingPullback.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/CopyableToDevice.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/DataTypes.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/DifferentialOperators.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/ElementaryFunctions.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/EuclideanDifferentiable.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Execution.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/KeyPathIterable.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorContext.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorOperation.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorShapeInference.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorTFFunctionBuilder.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorTrace.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/LazyTensorTraceCache.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/MixedPrecision.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/PointwiseMultiplicative.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/PythonConversion.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Random.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Runtime.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/ShapedArray.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/StringTensor.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Tensor.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/TensorGroup.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/TensorHandle.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/TensorShape.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/TensorUtilities.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Threading.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/Utilities.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Core/VectorProtocol.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/Algorithms.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/Backend.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/Collatable.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/NonuniformTrainingEpochs.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/Sampling.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/Slices.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Epochs/TrainingEpochs.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Freezable.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Initializers.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layer.swift -primary-file /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Convolutional.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Core.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Dense.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Dropout.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Embedding.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Initialization.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Morphological.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Normalization.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Pooling.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Recurrent.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Sequential.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layers/Upsampling.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Loss.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/Basic.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/Comparison.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/Image.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/LinearAlgebra.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/Math.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Operators/NN.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Optimizers/MomentumBased.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Optimizers/Optimizer.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Optimizers/SGD.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/StdlibExtensions.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/APIs/CrossReplicaSum.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/APIs/DeviceScope.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/APIs/RawOpsManual.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/Device.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/RawOpsXLAGenerated.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/XLAScalarType.swift /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/X10/XLATensor.swift -supplementary-output-file-map /var/folders/qn/86czb43d3pv03bfnxvb3x66h0000gn/T/TemporaryDirectory.aeNSps/supplementaryOutputs-6 -target arm64-apple-macosx10.13 -Xllvm -aarch64-use-tbi -enable-objc-interop -sdk /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.3.sdk -I /Users/philipturner/Documents/building-tensorflow/s4tf/.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/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/ModuleCache -swift-version 5 -Onone -D SWIFT_PACKAGE -D DEBUG -D DEFAULT_BACKEND_EAGER -new-driver-path /Library/Developer/Toolchains/swift-5.7-DEVELOPMENT-SNAPSHOT-2022-06-26-a.xctoolchain/usr/bin/swift-driver -empty-abi-descriptor -resource-dir /Library/Developer/Toolchains/swift-5.7-DEVELOPMENT-SNAPSHOT-2022-06-26-a.xctoolchain/usr/lib/swift -enable-anonymous-context-mangled-names -Xcc -fmodule-map-file=/Users/philipturner/Documents/building-tensorflow/s4tf/.build/checkouts/swift-numerics/Sources/_NumericsShims/include/module.modulemap -Xcc -I -Xcc /Users/philipturner/Documents/building-tensorflow/s4tf/.build/checkouts/swift-numerics/Sources/_NumericsShims/include -Xcc -fmodule-map-file=/Users/philipturner/Documents/building-tensorflow/s4tf/Sources/CX10Modules/include/module.modulemap -Xcc -I -Xcc /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/CX10Modules/include -Xcc -fmodule-map-file=/Users/philipturner/Documents/building-tensorflow/s4tf/Sources/CTensorFlow/include/module.modulemap -Xcc -I -Xcc /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/CTensorFlow/include -Xcc -I/Users/philipturner/Documents/building-tensorflow/Library/tensorflow-2.9.1/usr/include -module-name TensorFlow -target-sdk-version 12.3 -parse-as-library -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Epochs/NonuniformTrainingEpochs.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Epochs/Sampling.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Epochs/Slices.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Epochs/TrainingEpochs.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Freezable.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Initializers.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Layer.swift.o -o /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Layers/Convolutional.swift.o -index-store-path /Users/philipturner/Documents/building-tensorflow/s4tf/.build/arm64-apple-macosx/debug/index/store -index-system-modules
1.	Apple Swift version 5.7-dev (LLVM eecf02df5133efe, Swift ad06e7d7251311a)
2.	Compiling with the current language version
3.	While evaluating request ExecuteSILPipelineRequest(Run pipelines { Mandatory Diagnostic Passes + Enabling Optimization Passes } on SIL for TensorFlow)
4.	While running pass #1006 SILModuleTransform "Differentiation".
5.	While canonicalizing `differentiable_function` SIL node   %10 = differentiable_function [parameters 1] [results 0] %6 : $@callee_guaranteed (@in_guaranteed τ_0_0.Input, @in_guaranteed τ_0_0) -> @out τ_0_0.Output // users: %15, %11
6.	While ...in SIL function "@$s10TensorFlow5LayerPAAE7forwardy6OutputQz5InputQzFAaBRzlTJrUSpSr".
 for 'forward(_:)' (at /Users/philipturner/Documents/building-tensorflow/s4tf/Sources/TensorFlow/Layer.swift:170:10)
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           0x0000000104f83144 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
1  swift-frontend           0x0000000104f823a4 llvm::sys::RunSignalHandlers() + 128
2  swift-frontend           0x0000000104f837a8 SignalHandler(int) + 304
3  libsystem_platform.dylib 0x00000001877974a4 _sigtramp + 56
4  libsystem_pthread.dylib  0x000000018777fee0 pthread_kill + 288
5  libsystem_c.dylib        0x00000001876ba340 abort + 168
6  libsystem_c.dylib        0x00000001876b9754 err + 0
7  swift-frontend           0x0000000104fa2f04 swift::SILFunctionBuilder::getOrCreateFunction(swift::SILLocation, llvm::StringRef, swift::SILLinkage, swift::CanTypeWrapper<swift::SILFunctionType>, swift::IsBare_t, swift::IsTransparent_t, swift::IsSerialized_t, swift::IsDynamicallyReplaceable_t, swift::IsDistributed_t, swift::ProfileCounter, swift::IsThunk_t, swift::SubclassScope) (.cold.4) + 0
8  swift-frontend           0x0000000100b42ef4 swift::SILFunctionBuilder::getOrCreateFunction(swift::SILLocation, llvm::StringRef, swift::SILLinkage, swift::CanTypeWrapper<swift::SILFunctionType>, swift::IsBare_t, swift::IsTransparent_t, swift::IsSerialized_t, swift::IsDynamicallyReplaceable_t, swift::IsDistributed_t, swift::ProfileCounter, swift::IsThunk_t, swift::SubclassScope) + 512
9  swift-frontend           0x0000000100b4425c swift::SILFunctionBuilder::getOrCreateSharedFunction(swift::SILLocation, llvm::StringRef, swift::CanTypeWrapper<swift::SILFunctionType>, swift::IsBare_t, swift::IsTransparent_t, swift::IsSerialized_t, swift::ProfileCounter, swift::IsThunk_t, swift::IsDynamicallyReplaceable_t, swift::IsDistributed_t) + 76
10 swift-frontend           0x000000010112f3dc swift::autodiff::getOrCreateSubsetParametersThunkForLinearMap(swift::SILOptFunctionBuilder&, swift::SILFunction*, swift::CanTypeWrapper<swift::SILFunctionType>, swift::CanTypeWrapper<swift::SILFunctionType>, swift::CanTypeWrapper<swift::SILFunctionType>, swift::AutoDiffDerivativeFunctionKind, swift::AutoDiffConfig const&, swift::AutoDiffConfig const&, swift::autodiff::ADContext&) + 976
11 swift-frontend           0x00000001011311c0 swift::autodiff::getOrCreateSubsetParametersThunkForDerivativeFunction(swift::SILOptFunctionBuilder&, swift::SILValue, swift::SILValue, swift::AutoDiffDerivativeFunctionKind, swift::AutoDiffConfig const&, swift::AutoDiffConfig const&, swift::autodiff::ADContext&) + 2880
12 swift-frontend           0x0000000101237db0 (anonymous namespace)::DifferentiationTransformer::promoteToDifferentiableFunction(swift::DifferentiableFunctionInst*, swift::SILBuilder&, swift::SILLocation, swift::autodiff::DifferentiationInvoker) + 6784
13 swift-frontend           0x0000000101234814 (anonymous namespace)::DifferentiationTransformer::processDifferentiableFunctionInst(swift::DifferentiableFunctionInst*) + 484
14 swift-frontend           0x0000000101232898 (anonymous namespace)::Differentiation::run() + 1292
15 swift-frontend           0x00000001012ab058 swift::SILPassManager::runModulePass(unsigned int) + 956
16 swift-frontend           0x00000001012b05fc swift::SILPassManager::execute() + 624
17 swift-frontend           0x00000001012a8138 swift::SILPassManager::executePassPipelinePlan(swift::SILPassPipelinePlan const&) + 68
18 swift-frontend           0x00000001012a80c0 swift::ExecuteSILPipelineRequest::evaluate(swift::Evaluator&, swift::SILPipelineExecutionDescriptor) const + 68
19 swift-frontend           0x00000001012c86f0 swift::SimpleRequest<swift::ExecuteSILPipelineRequest, std::__1::tuple<> (swift::SILPipelineExecutionDescriptor), (swift::RequestFlags)1>::evaluateRequest(swift::ExecuteSILPipelineRequest const&, swift::Evaluator&) + 28
20 swift-frontend           0x00000001012b31b8 llvm::Expected<swift::ExecuteSILPipelineRequest::OutputType> swift::Evaluator::getResultUncached<swift::ExecuteSILPipelineRequest>(swift::ExecuteSILPipelineRequest const&) + 252
21 swift-frontend           0x00000001012a8334 swift::executePassPipelinePlan(swift::SILModule*, swift::SILPassPipelinePlan const&, bool, swift::irgen::IRGenModule*) + 84
22 swift-frontend           0x00000001012b6cbc swift::runSILDiagnosticPasses(swift::SILModule&) + 92
23 swift-frontend           0x0000000100af2868 swift::CompilerInstance::performSILProcessing(swift::SILModule*) + 72
24 swift-frontend           0x0000000100a9740c 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*) + 716
25 swift-frontend           0x0000000100a96be8 swift::performCompileStepsPostSema(swift::CompilerInstance&, int&, swift::FrontendObserver*) + 540
26 swift-frontend           0x0000000100aa4a34 withSemanticAnalysis(swift::CompilerInstance&, swift::FrontendObserver*, llvm::function_ref<bool (swift::CompilerInstance&)>, bool) + 160
27 swift-frontend           0x0000000100a98a78 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 3388
28 swift-frontend           0x000000010099f08c swift::mainEntry(int, char const**) + 3764
29 dyld                     0x000000010d8b908c start + 520
^C

Environment (please fill out the following information)

• OS: macOS 12.4
• Xcode Version/Tag/Branch: 2022-06-26 v5.7 Development Snapshot, Xcode 14 beta 2 Swift 5.7 Release Toolchain

Additional context
Neither of my toolchains have recent cross-file derivative registration fixes in yet because of #59467. The crasher could have disappeared since the patch.