[AutoDiff] Assertion failed: `(linearMapStructTy && "must have linear map struct type for predecessor BB")`

[fibrechannelscsi]

Description

The code below will cause the compiler to crash when compiled with certain toolchains. The error message provided is:
Assertion failed: (linearMapStructTy && "must have linear map struct type for predecessor BB"), function populateBranchingTraceDecl, file LinearMapInfo.cpp, line 146.

The error occurs in both Debug and Release mode.

Reproduction

Compile the following code in Debug or Release mode:

import Foundation; import _Differentiation
struct A: Differentiable {}
struct B: Differentiable {@differentiable(reverse) func c(b: B) -> A {while true {if true {break}}; return A()}}

Expected behavior

The compilation should succeed.

Environment

This fails with:
Swift Development Snapshot 2024-01-15(a)
Swift Development Snapshot 2024-01-22(a)

but not with:
Swift Development Snapshot 2024-01-08(a)
Swift 5.10 Development Snapshot 2024-01-24(a)

This failure can be seen on both an M1 and x86 Mac.

Additional information

Here is a stack trace:

1.	Apple Swift version 5.11-dev (LLVM e2b71339d3295a4, Swift 8ae3ee998f408c4)
2.	Compiling with the current language version
3.	While evaluating request ExecuteSILPipelineRequest(Run pipelines { Mandatory Diagnostic Passes + Enabling Optimization Passes } on SIL for linearMapReproducer)
4.	While running pass #388 SILModuleTransform "Differentiation".
5.	While processing // differentiability witness for B.c(b:)
sil_differentiability_witness [serialized] [reverse] [parameters 0 1] [results 0] @$s19linearMapReproducer1BV1c1bAA1AVAC_tF : $@convention(method) (B, B) -> A {
}

 on SIL function "@$s19linearMapReproducer1BV1c1bAA1AVAC_tF".
 for 'c(b:)' (at /Users/user/linearMapReproducer/linearMapReproducer/main.swift:3:52)
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           0x0000000107ed82c8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
1  swift-frontend           0x0000000107ed6a94 llvm::sys::RunSignalHandlers() + 112
2  swift-frontend           0x0000000107ed8910 SignalHandler(int) + 304
3  libsystem_platform.dylib 0x000000018b12aa24 _sigtramp + 56
4  libsystem_pthread.dylib  0x000000018b0fbc28 pthread_kill + 288
5  libsystem_c.dylib        0x000000018b009ae8 abort + 180
6  libsystem_c.dylib        0x000000018b008e44 err + 0
7  swift-frontend           0x000000010803de88 swift::autodiff::LinearMapInfo::populateBranchingTraceDecl(swift::SILBasicBlock*, swift::SILLoopInfo*) (.cold.4) + 0
8  swift-frontend           0x0000000103258d18 swift::autodiff::LinearMapInfo::populateBranchingTraceDecl(swift::SILBasicBlock*, swift::SILLoopInfo*) + 1196
9  swift-frontend           0x000000010325799c swift::autodiff::LinearMapInfo::generateDifferentiationDataStructures(swift::autodiff::ADContext&, swift::SILFunction*) + 452
10 swift-frontend           0x0000000103258160 swift::autodiff::LinearMapInfo::LinearMapInfo(swift::autodiff::ADContext&, swift::AutoDiffLinearMapKind, swift::SILFunction*, swift::SILFunction*, swift::AutoDiffConfig const&, swift::DifferentiableActivityInfo const&, swift::SILLoopInfo*) + 148
11 swift-frontend           0x0000000103277f68 swift::autodiff::VJPCloner::Implementation::Implementation(swift::autodiff::VJPCloner&, swift::autodiff::ADContext&, swift::SILDifferentiabilityWitness*, swift::SILFunction*, swift::autodiff::DifferentiationInvoker) + 656
12 swift-frontend           0x0000000103278e2c swift::autodiff::VJPCloner::VJPCloner(swift::autodiff::ADContext&, swift::SILDifferentiabilityWitness*, swift::SILFunction*, swift::autodiff::DifferentiationInvoker) + 92
13 swift-frontend           0x00000001033d4bec (anonymous namespace)::DifferentiationTransformer::canonicalizeDifferentiabilityWitness(swift::SILDifferentiabilityWitness*, swift::autodiff::DifferentiationInvoker, swift::IsSerialized_t) + 6032
14 swift-frontend           0x00000001033d2a64 (anonymous namespace)::Differentiation::run() + 1060
15 swift-frontend           0x0000000103470f64 swift::SILPassManager::runModulePass(unsigned int) + 848
16 swift-frontend           0x0000000103473198 swift::SILPassManager::execute() + 624
17 swift-frontend           0x000000010346dbcc swift::SILPassManager::executePassPipelinePlan(swift::SILPassPipelinePlan const&) + 72
18 swift-frontend           0x000000010346db50 swift::ExecuteSILPipelineRequest::evaluate(swift::Evaluator&, swift::SILPipelineExecutionDescriptor) const + 68
19 swift-frontend           0x00000001034a7cf0 swift::SimpleRequest<swift::ExecuteSILPipelineRequest, std::__1::tuple<> (swift::SILPipelineExecutionDescriptor), (swift::RequestFlags)1>::evaluateRequest(swift::ExecuteSILPipelineRequest const&, swift::Evaluator&) + 28
20 swift-frontend           0x0000000103488cf4 llvm::Expected<swift::ExecuteSILPipelineRequest::OutputType> swift::Evaluator::getResultUncached<swift::ExecuteSILPipelineRequest>(swift::ExecuteSILPipelineRequest const&) + 268
21 swift-frontend           0x000000010346ddb8 swift::executePassPipelinePlan(swift::SILModule*, swift::SILPassPipelinePlan const&, bool, swift::irgen::IRGenModule*) + 84
22 swift-frontend           0x000000010348b318 swift::runSILDiagnosticPasses(swift::SILModule&) + 192
23 swift-frontend           0x0000000102d0560c swift::CompilerInstance::performSILProcessing(swift::SILModule*) + 80
24 swift-frontend           0x0000000102ad5e44 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*) + 796
25 swift-frontend           0x0000000102ad5538 swift::performCompileStepsPostSema(swift::CompilerInstance&, int&, swift::FrontendObserver*) + 624
26 swift-frontend           0x0000000102ae5ff8 withSemanticAnalysis(swift::CompilerInstance&, swift::FrontendObserver*, llvm::function_ref<bool (swift::CompilerInstance&)>, bool) + 160
27 swift-frontend           0x0000000102ad7a48 performCompile(swift::CompilerInstance&, int&, swift::FrontendObserver*) + 708
28 swift-frontend           0x0000000102ad6a7c swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 2280
29 swift-frontend           0x000000010295944c swift::mainEntry(int, char const**) + 3064
30 dyld                     0x000000018ada3f28 start + 2236

[fibrechannelscsi]

@asl

[asl]

@fibrechannelscsi I cannot reproduce this on main. Will it be possible for you to check it again? (in the meantime, rechecking with 8ae3ee9)

[fibrechannelscsi]

I've checked this with the most recent 01-30 nightlies:
https://download.swift.org/development/xcode/swift-DEVELOPMENT-SNAPSHOT-2024-01-30-a/swift-DEVELOPMENT-SNAPSHOT-2024-01-30-a-osx.pkg
and
https://download.swift.org/development/ubuntu1804/swift-DEVELOPMENT-SNAPSHOT-2024-01-30-a/swift-DEVELOPMENT-SNAPSHOT-2024-01-30-a-ubuntu18.04.tar.gz
in the case of an x86 Ubuntu 18 machine.

This is still occurring both on x86 and on my M1 Mac with the 01-30a toolchains.
If I copy the three lines above into a main.swift file and simply run:

/Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2024-01-30-a.xctoolchain/usr/bin/swift main.swift
then it will segfault as before. Given that I also have the 01-08a toolchain installed:
/Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2024-01-08-a.xctoolchain/usr/bin/swift main.swift
will exit without any output as expected.

[asl]

Interesting, thanks!

[asl]

Ok, I was able to reproduce in bootstrapped build (not the build with bootstrapping turned off)! @eeckstein Is this something expected? This is w/o any optimizations...

[asl]

Here is input SIL for function B:

// B.c(b:)
sil hidden [ossa] @$s4main1BV1c1bAA1AVAC_tF : $@convention(method) (B, B) -> A {
// %0 "b"                                         // user: %2
// %1 "self"                                      // user: %3
bb0(%0 : $B, %1 : $B):
  debug_value %0 : $B, let, name "b", argno 1     // id: %2
  debug_value %1 : $B, let, name "self", argno 2, implicit // id: %3
  br bb1                                          // id: %4

bb1:                                              // Preds: bb4 bb0
  %5 = integer_literal $Builtin.Int1, -1          // user: %8
  %6 = metatype $@thin Bool.Type                  // user: %8
  // function_ref Bool.init(_builtinBooleanLiteral:)
  %7 = function_ref @$sSb22_builtinBooleanLiteralSbBi1__tcfC : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %8
  %8 = apply %7(%5, %6) : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %9
  %9 = struct_extract %8 : $Bool, #Bool._value
  br bb2                                          // id: %10

bb2:                                              // Preds: bb1
  %11 = integer_literal $Builtin.Int1, -1         // user: %14
  %12 = metatype $@thin Bool.Type                 // user: %14
  // function_ref Bool.init(_builtinBooleanLiteral:)
  %13 = function_ref @$sSb22_builtinBooleanLiteralSbBi1__tcfC : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %14
  %14 = apply %13(%11, %12) : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %15
  %15 = struct_extract %14 : $Bool, #Bool._value
  br bb3                                          // id: %16

bb3:                                              // Preds: bb2
  br bb6                                          // id: %17

bb4:
  br bb1                                          // id: %18

bb5:
  br bb6                                          // id: %19

bb6:                                              // Preds: bb3 bb5
  %20 = metatype $@thin A.Type                    // user: %22
  // function_ref A.init()
  %21 = function_ref @$s4main1AVACycfC : $@convention(method) (@thin A.Type) -> A // user: %22
  %22 = apply %21(%20) : $@convention(method) (@thin A.Type) -> A // user: %23
  return %22 : $A                                 // id: %23
} // end sil function '$s4main1BV1c1bAA1AVAC_tF'

bb4 here is unreachable and therefore excluded from RPOT. Here is the input SIL without bootstrap (so with all swift code in compiler disabled):

sil hidden [ossa] @$s4main1BV1c1bAA1AVAC_tF : $@convention(method) (B, B) -> A {
// %0 "b"                                         // user: %2
// %1 "self"                                      // user: %3
bb0(%0 : $B, %1 : $B):
  debug_value %0 : $B, let, name "b", argno 1     // id: %2
  debug_value %1 : $B, let, name "self", argno 2, implicit // id: %3
  br bb1                                          // id: %4

bb1:                                              // Preds: bb4 bb0
  %5 = integer_literal $Builtin.Int1, -1          // user: %8
  %6 = metatype $@thin Bool.Type                  // user: %8
  // function_ref Bool.init(_builtinBooleanLiteral:)
  %7 = function_ref @$sSb22_builtinBooleanLiteralSbBi1__tcfC : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %8
  %8 = apply %7(%5, %6) : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %9
  %9 = struct_extract %8 : $Bool, #Bool._value    // user: %10
  cond_br %9, bb2, bb5                            // id: %10

bb2:                                              // Preds: bb1
  %11 = integer_literal $Builtin.Int1, -1         // user: %14
  %12 = metatype $@thin Bool.Type                 // user: %14
  // function_ref Bool.init(_builtinBooleanLiteral:)
  %13 = function_ref @$sSb22_builtinBooleanLiteralSbBi1__tcfC : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %14
  %14 = apply %13(%11, %12) : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool // user: %15
  %15 = struct_extract %14 : $Bool, #Bool._value  // user: %16
  cond_br %15, bb3, bb4                           // id: %16

bb3:                                              // Preds: bb2
  br bb6                                          // id: %17

bb4:                                              // Preds: bb2
  br bb1                                          // id: %18

bb5:                                              // Preds: bb1
  br bb6                                          // id: %19

bb6:                                              // Preds: bb3 bb5
  %20 = metatype $@thin A.Type                    // user: %22
  // function_ref A.init()
  %21 = function_ref @$s4main1AVACycfC : $@convention(method) (@thin A.Type) -> A // user: %22
  %22 = apply %21(%20) : $@convention(method) (@thin A.Type) -> A // user: %23
  return %22 : $A                                 // id: %23
} // end sil function '$s4main1BV1c1bAA1AVAC_tF'

Clearly, some pass replaced cond_br %15, bb3, bb4 to direct branch making bb4 unreachable.

[asl]

It looks like it should be #70787

Tagging @eeckstein

[asl]

Some further information: for autodiff we'd essentially need to execute code in reverse from result to entry basic block. Therefore such unreachable blocks become reachable (on CFG with all edges reversed).

So I guess we'd need to detect such cases consistently and ignore blocks that are not reachable from the input block. Alternatively, we can just remove all unreachable blocks, this might be a cleaner solution actually.

[eeckstein]

@asl Is this something you can fix in the autodiff code? (I have no idea how that works)

[asl]

@asl Is this something you can fix in the autodiff code? (I have no idea how that works)

I have a fix for autodiff, albeit I'm not 100% sure it covers all cases :) I guess here is the question about assumptions and invariants: so far there were no unreachable BBs created during mandatory transformations. Will it make sense to remove them instead?