[AutoDiff][IRGen] Can't compile Swift for TensorFlow on top-of-tree Swift

[philipturner]

Describe the bug
This bug appeared between the May 11, 2022 and May 18, 2022 Trunk Development Snapshots. It seems to be caused by changing behavior of cross-file lookup for derivatives, resulting from the merging of #58644. It seems similar to #55170. The suspected source of this bug was not merged into the release/5.7 branch, so it is not a great concern for now. I will still be able to compile on the Swift 5.7 release toolchain, which I am currently targeting.

To Reproduce
Steps to reproduce the behavior:

1. Make a Swift package with the following Package.swift:

import PackageDescription

let package = Package(
  name: "TensorFlow",
  products: [
    .library(
      name: "TensorFlow",
      targets: ["TensorFlow"]),
  ],
  dependencies: [],
  targets: [
    .target(
      name: "TensorFlow",
      dependencies: []),
  ]
)

2. The following under Sources/TensorFlow/Core/Tensor.swift:

import _Differentiation

struct Tensor: Differentiable {}

// `Tensor` could be defined in "Normalization.swift" and the crash would still
// happen. All that matters is that `LayerNorm_callAsFunction` and `rsqrt` are
// defined in separate files.

@differentiable(reverse)
func rsqrt( _ x: Tensor) -> Tensor {
  fatalError()
}

@derivative(of: rsqrt)
func _vjpRsqrt(_ x: Tensor) -> (
  value: Tensor, pullback: (Tensor.TangentVector) -> (Tensor.TangentVector)
) {
  fatalError()
}

3. The following under Sources/TensorFlow/Core/Layers/Normalization.swift

import _Differentiation

@_semantics("autodiff.nonvarying")
func withoutDerivative() -> Tensor {
  fatalError()
}

func BatchNorm_doInference(
  _ input: Tensor
) -> Tensor {
  withoutDerivative()
}

@differentiable(reverse)
func BatchNorm_callAsFunction(_ input: Tensor) -> Tensor {
  BatchNorm_doInference(input)
}

@differentiable(reverse)
func LayerNorm_callAsFunction(_ input: Tensor) -> Tensor {
  rsqrt(input)
}

Crash stack trace

Failed to reconstruct type for $s10TensorFlow09_AD__$s10A58Flow21BatchNorm_doInferenceyAA0A0VADF_bb0__PB__src_0_wrt_0VmD
Original type:
(metatype_type
  (struct_type decl=TensorFlow.(file)._AD__$s10TensorFlow21BatchNorm_doInferenceyAA0A0VADF_bb0__PB__src_0_wrt_0))
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-DEVELOPMENT-SNAPSHOT-2022-06-02-a.xctoolchain/usr/bin/swift-frontend -frontend -c /Users/philipturner/Desktop/fan/s4tf/Sources/TensorFlow/Core/Tensor.swift -primary-file /Users/philipturner/Desktop/fan/s4tf/Sources/TensorFlow/Layers/Normalization.swift -emit-dependencies-path /Users/philipturner/Desktop/fan/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Layers/Normalization.d -emit-reference-dependencies-path /Users/philipturner/Desktop/fan/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Layers/Normalization.swiftdeps -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/fan/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/Desktop/fan/s4tf/.build/arm64-apple-macosx/debug/ModuleCache -swift-version 5 -Onone -D SWIFT_PACKAGE -D DEBUG -new-driver-path /Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2022-06-02-a.xctoolchain/usr/bin/swift-driver -empty-abi-descriptor -resource-dir /Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2022-06-02-a.xctoolchain/usr/lib/swift -enable-anonymous-context-mangled-names -module-name TensorFlow -target-sdk-version 12.3 -parse-as-library -o /Users/philipturner/Desktop/fan/s4tf/.build/arm64-apple-macosx/debug/TensorFlow.build/Layers/Normalization.swift.o -index-store-path /Users/philipturner/Desktop/fan/s4tf/.build/arm64-apple-macosx/debug/index/store -index-system-modules
1.	Apple Swift version 5.8-dev (LLVM 278d67f38c6a910, Swift ee312bc1e20eb01)
2.	Compiling with the current language version
3.	While evaluating request IRGenRequest(IR Generation for file "/Users/philipturner/Desktop/fan/s4tf/Sources/TensorFlow/Layers/Normalization.swift")
4.	While emitting IR for synthesized file0x11a1d01d8
5.	While emitting metadata for '_AD__$s10TensorFlow21BatchNorm_doInferenceyAA0A0VADF_bb0__PB__src_0_wrt_0' (in module 'TensorFlow')
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           0x0000000106b823a4 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
1  swift-frontend           0x0000000106b81604 llvm::sys::RunSignalHandlers() + 128
2  swift-frontend           0x0000000106b82a08 SignalHandler(int) + 304
3  libsystem_platform.dylib 0x00000001a6a674a4 _sigtramp + 56
4  libsystem_pthread.dylib  0x00000001a6a4fee0 pthread_kill + 288
5  libsystem_c.dylib        0x00000001a698a340 abort + 168
6  swift-frontend           0x0000000106c5eed4 (anonymous namespace)::IRGenDebugInfoImpl::getOrCreateType(swift::irgen::DebugTypeInfo) (.cold.9) + 0
7  swift-frontend           0x0000000102a59950 (anonymous namespace)::IRGenDebugInfoImpl::getOrCreateType(swift::irgen::DebugTypeInfo) + 3788
8  swift-frontend           0x0000000102a54e80 swift::irgen::IRGenDebugInfo::emitGlobalVariableDeclaration(llvm::GlobalVariable*, llvm::StringRef, llvm::StringRef, swift::irgen::DebugTypeInfo, bool, bool, llvm::Optional<swift::SILLocation>) + 540
9  swift-frontend           0x0000000102969660 swift::irgen::createVariable(swift::irgen::IRGenModule&, swift::irgen::LinkInfo&, llvm::Type*, swift::irgen::Alignment, swift::irgen::DebugTypeInfo, llvm::Optional<swift::SILLocation>, llvm::StringRef, bool) + 584
10 swift-frontend           0x000000010296dc90 swift::irgen::IRGenModule::getAddrOfLLVMVariable(swift::irgen::LinkEntity, swift::irgen::ConstantInit, swift::irgen::DebugTypeInfo, llvm::Type*) + 776
11 swift-frontend           0x00000001029707e0 swift::irgen::IRGenModule::defineTypeMetadata(swift::CanType, bool, bool, clang::CodeGen::ConstantInitFuture, llvm::StringRef, llvm::SmallVector<std::__1::pair<swift::irgen::Size, swift::SILDeclRef>, 8u>) + 468
12 swift-frontend           0x00000001029d6a58 swift::irgen::emitStructMetadata(swift::irgen::IRGenModule&, swift::StructDecl*) + 3872
13 swift-frontend           0x0000000102a0fa54 swift::irgen::IRGenModule::emitStructDecl(swift::StructDecl*) + 56
14 swift-frontend           0x000000010295f880 swift::irgen::IRGenModule::emitSynthesizedFileUnit(swift::SynthesizedFileUnit&) + 116
15 swift-frontend           0x0000000102a41ba0 swift::IRGenRequest::evaluate(swift::Evaluator&, swift::IRGenDescriptor) const + 1996
16 swift-frontend           0x0000000102a78cb0 swift::GeneratedModule swift::SimpleRequest<swift::IRGenRequest, swift::GeneratedModule (swift::IRGenDescriptor), (swift::RequestFlags)9>::callDerived<0ul>(swift::Evaluator&, std::__1::integer_sequence<unsigned long, 0ul>) const + 196
17 swift-frontend           0x0000000102a4c9d8 llvm::Expected<swift::IRGenRequest::OutputType> swift::Evaluator::getResultUncached<swift::IRGenRequest>(swift::IRGenRequest const&) + 372
18 swift-frontend           0x0000000102a43dec 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**) + 212
19 swift-frontend           0x00000001026c450c 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*) + 1664
20 swift-frontend           0x00000001026c3930 swift::performCompileStepsPostSema(swift::CompilerInstance&, int&, swift::FrontendObserver*) + 540
21 swift-frontend           0x00000001026d1584 withSemanticAnalysis(swift::CompilerInstance&, swift::FrontendObserver*, llvm::function_ref<bool (swift::CompilerInstance&)>, bool) + 160
22 swift-frontend           0x00000001026c5660 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 3388
23 swift-frontend           0x00000001025ace44 swift::mainEntry(int, char const**) + 3764
24 dyld                     0x000000010f58508c start + 520

Environment (please complete the following information):

• OS: Ubuntu 18.04, x86_64 (Google Colab)
• Toolchain: First appeared on the May 18, 2022 Trunk Development Snapshot

[philipturner]

CC: @asl hopefully #58965 resolves this bug.

[asl]

Seems to be unrelated to #58965 (and certainly not similar to #55170) as assertion is in debug info generation. Likely compilation w/o debug info will allow to workaround an issue

[BradLarson]

@asl - While it points to debug info, we had a very similar "Failed to reconstruct type" error appear in a local module that did seem to go away after #58965 was applied. This is a much simpler reproducer than our case, but the trace looks close.

[philipturner]

I found the reproducer while testing an old branch that was usable only months ago. It was pretty much the same as Fan's branch, which was only usable with the November 12, 2021 toolchain. But I just tested top-of-tree S4TF at s4tf/s4tf, and the crash appears there as well. I'm surprised it slipped under my nose for almost 30 days.

I had been testing only the 5.7 branch snapshots since exactly May 18. I had both the trunk toolchain and 5.7 branch toolchain for May 18 on my computer, and thought the trunk toolchain was running. Apparently it was the 5.7 branch toolchain.

https://gist.github.com/philipturner/7aa063af04277d463c14168275878511/revisions#diff-5120188d4307c1971360f954563798c33e9af6966e593cb90ec127f6a3d5c3bbR123

https://gist.github.com/philipturner/7aa063af04277d463c14168275878511/revisions#diff-5120188d4307c1971360f954563798c33e9af6966e593cb90ec127f6a3d5c3bbR151

[philipturner]

I am unable to reproduce the crash on top-of-tree as a Stdlib regression test. I tried replication the style of @asl's test/AutoDiff/Sema/DerivativeRegistrationCrossFile, but no luck. I also tried running the test on the #58965 branch, and nothing changed. To add this bug as a regression test for #58965, we need to investigate it further.

The base Swift repository doesn't include swift-package, so we can't build via the Swift package manager. I have not succeeded in reproducing the crash on the 2022-06-18 toolchain using the swiftc command; rather, I have to use SwiftPM. That means we must first discover how the SwiftPM build translates into one massive swiftc call. Once that is discovered, we can add this as a test case.

Here is the progress I made. Add this file to test/AutoDiff/compiler_crashers_fixed/Inputs as 59467-???.swift:

import _Differentiation

struct Tensor: Differentiable {}

// `Tensor` could be defined in this test case's primary file and the crash
// would still happen. All that matters is that `LayerNorm_callAsFunction` and
// `rsqrt` are defined in separate files.

@differentiable(reverse)
func rsqrt( _ x: Tensor) -> Tensor {
  fatalError()
}

@derivative(of: rsqrt)
func _vjpRsqrt(_ x: Tensor) -> (
  value: Tensor, pullback: (Tensor.TangentVector) -> (Tensor.TangentVector)
) {
  fatalError()
}

And add the following as 59467-???.swift. You can choose the file name that seems most appropriate, replacing "???" with the chosen name.

// RUN: %target-swift-frontend -emit-ir -primary-file %s %S/Inputs/59467-???.swift -module-name main -o /dev/null

import _Differentiation

@_semantics("autodiff.nonvarying")
func withoutDerivative() -> Tensor {
  fatalError()
}

func BatchNorm_doInference(
  _ input: Tensor
) -> Tensor {
  withoutDerivative()
}

@differentiable(reverse)
func BatchNorm_callAsFunction(_ input: Tensor) -> Tensor {
  BatchNorm_doInference(input)
}

@differentiable(reverse)
func LayerNorm_callAsFunction(_ input: Tensor) -> Tensor {
  rsqrt(input)
}

Would you be able to finish what I started here?

[asl]

We should not add this as #58965 regression test. As the issue is completely different. It might be hidden by #58965 actually.

I also tried running the test on the #58965 branch, and nothing changed.

What do you mean? The issue was resolved? Or not resolved?

[philipturner]

The compiler failed to crash, just like it failed to crash on the main branch. We don't yet know whether #58965 fixed, hid, or did nothing to this bug. We don't yet have a way to test #58965 against the bug. That's what I was getting at.

[asl]

Can't you simply build a toolchain from branch and build your code with it?

[philipturner]

The closest I have to building a toolchain from scratch with my current skillset is:

1. Compile the Swift compiler using the instructions on GettingStarted.md.
2. Scour the build products for the swiftc binary and call that binary from the command line.

I don't know how to compile a full-on toolchain that includes SwiftPM and everything added on to the bare apple/swift repo. I was hoping that you could do so more easily than I could.

[asl]

Hopefully this: https://github.com/apple/swift-package-manager/blob/main/CONTRIBUTING.md will help. Check Advanced section.

[philipturner]

I'm still struggling to get that working. Overall, I think it's more time-effective if you test this bug with SwiftPM. Plus, it was PassiveLogic that originally encountered the bug (#59467 (comment)) without doing the necessary work to fully narrow down and report it. I don't mean to be disrespectful, but I have been narrowing this bug for ~12 hours straight and I think it logically makes sense that you and Brad share some of the responsibility of investigating this bug.

[asl]

I'll let @BradLarson and PassiveLogic folks take care about test reduction for you.