[AutoDiff] Conflicting debug info for argument with custom valueWithPullback

[BradLarson]

Description

A specific configuration of a wrapper function around valueWithPullback() built with both optimization and debug symbol generation leads to an assertion failure of conflicting debug info for argument.

Steps to reproduce

Place the following in a file:

import _Differentiation

public extension Array {
    @inlinable
    @differentiable(reverse)
    mutating func update(at index: Int, byCalling closure: @differentiable(reverse) (inout Element) -> Void) where Element: Differentiable {
        closure(&self[index])
    }
}

public func valueWithPullback<T>(
    at x: T, of f: @differentiable(reverse) (inout T) -> Void
) -> (value: Void, pullback: (inout T.TangentVector) -> Void) {
    @differentiable(reverse)
    func nonInoutWrappingFunction(_ t: T) -> T {
        var t = t
        f(&t)
        return t
    }
    let nonInoutPullback = pullback(at: x, of: nonInoutWrappingFunction)
    return ((), { $0 = nonInoutPullback($0) })
}

@inlinable
public func pullback<T>(
    at x: T, of f: @differentiable(reverse) (inout T) -> Void
) -> (inout T.TangentVector) -> Void {
    return valueWithPullback(at: x, of: f).pullback
}

public extension Array where Element: Differentiable {
    @inlinable
    @derivative(of: update(at:byCalling:))
    mutating func vjpUpdate(
        at index: Int,
        byCalling closure: @differentiable(reverse) (inout Element) -> Void
    )
    ->
    (value: Void, pullback: (inout Self.TangentVector) -> Void)
    {
        let closurePullback = pullback(at: self[index], of: closure)
        return (value: (), pullback: { closurePullback(&$0.base[index]) })
    }
}

func testUpdateByCalling() throws {
    @differentiable(reverse)
    func fOfArray(array: [Double]) -> Double {
        var array = array
        var result = 0.0
        for i in withoutDerivative(at: 0 ..< array.count) {
            array.update(at: i, byCalling: { (element: inout Double) in
                let initialElement = element
                for _ in withoutDerivative(at: 0 ..< i) {
                    element *= initialElement
                }
            })
            result += array[i]
        }
        return result
    }
    
    let array = [Double](repeating: 1.0, count: 3)
    let expectedGradientOfFOfArray = [1.0, 2.0, 3.0]
    let obtainedGradientOfFOfArray = gradient(at: array, of: fOfArray).base
}

And build it via swiftc -O -g file.swift. Note that both optimization and debug symbol generation must be used for this assertion failure to happen.

The assertion failure looks something like the following:

conflicting debug info for argument
  call void @llvm.dbg.value(metadata double %14, metadata !797, metadata !DIExpression()), !dbg !807
!796 = !DILocalVariable(name: "element", arg: 1, scope: !792, file: !18, line: 53, type: !201)
!797 = !DILocalVariable(name: "element", arg: 1, scope: !792, file: !18, line: 53, type: !201, flags: DIFlagArtificial)
<unknown>:0: error: fatal error encountered during compilation; please submit a bug report (https://swift.org/contributing/#reporting-bugs) and include the project
<unknown>:0: note: Broken module found, compilation aborted!
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-10-09-a.xctoolchain/usr/bin/swift-frontend -frontend -c -primary-file ConflictingDebugInfo.swift -target arm64-apple-macosx12.0 -Xllvm -aarch64-use-tbi -enable-objc-interop -sdk /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk -color-diagnostics -g -O -new-driver-path /Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2022-10-09-a.xctoolchain/usr/bin/swift-driver -empty-abi-descriptor -resource-dir /Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2022-10-09-a.xctoolchain/usr/lib/swift -module-name ConflictingDebugInfo -target-sdk-version 13.0 -target-sdk-name macosx13.0 -enable-default-cmo -o /var/folders/6y/ghmhq_xj3xq100vsttqqn2lc0000gq/T/TemporaryDirectory.3eixJp/ConflictingDebugInfo-2.o
1.	Apple Swift version 5.8-dev (LLVM b9ce35da9a19288, Swift 7335d60c50dc6af)
2.	Compiling with the current language version
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           0x0000000107224840 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
1  swift-frontend           0x0000000107223ac4 llvm::sys::RunSignalHandlers() + 128
2  swift-frontend           0x0000000107224e80 SignalHandler(int) + 304
3  libsystem_platform.dylib 0x00000001b77c74a4 _sigtramp + 56
4  libsystem_pthread.dylib  0x00000001b77afee0 pthread_kill + 288
5  libsystem_c.dylib        0x00000001b76ea340 abort + 168
6  swift-frontend           0x000000010295a924 PrettyStackTraceFrontend::~PrettyStackTraceFrontend() + 0
7  swift-frontend           0x0000000107188204 llvm::report_fatal_error(llvm::Twine const&, bool) + 268
8  swift-frontend           0x00000001071880f8 llvm::report_fatal_error(llvm::Twine const&, bool) + 0
9  swift-frontend           0x0000000106fcfaa8 llvm::VerifierPass::run(llvm::Function&, llvm::AnalysisManager<llvm::Function>&) + 0
10 swift-frontend           0x0000000106fad380 llvm::PassManager<llvm::Module, llvm::AnalysisManager<llvm::Module> >::run(llvm::Module&, llvm::AnalysisManager<llvm::Module>&) + 396
11 swift-frontend           0x0000000102d09ecc swift::performLLVMOptimizations(swift::IRGenOptions const&, llvm::Module*, llvm::TargetMachine*) + 4012
12 swift-frontend           0x0000000102d0b0c4 swift::performLLVM(swift::IRGenOptions const&, swift::DiagnosticEngine&, llvm::sys::SmartMutex<false>*, llvm::GlobalVariable*, llvm::Module*, llvm::TargetMachine*, llvm::StringRef, swift::UnifiedStatsReporter*) + 3092
13 swift-frontend           0x0000000102954a74 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*) + 2512
14 swift-frontend           0x0000000102953afc swift::performCompileStepsPostSema(swift::CompilerInstance&, int&, swift::FrontendObserver*) + 576
15 swift-frontend           0x00000001029628b0 withSemanticAnalysis(swift::CompilerInstance&, swift::FrontendObserver*, llvm::function_ref<bool (swift::CompilerInstance&)>, bool) + 160
16 swift-frontend           0x0000000102955824 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 3256
17 swift-frontend           0x00000001027b014c swift::mainEntry(int, char const**) + 3288
18 dyld                     0x000000011075908c start + 520

Expected behavior

The file should build cleanly.

Environment

• Swift compiler version: present in current main, and extending back in nightly snapshots for as far as I've tried
• Deployment target: tested on macOS 12.5.1

[BradLarson]

cc @asl

[asl]

The offending function is $s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr. Here we're having:

define internal swiftcc void @"$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr"(%TSd* nocapture dereferenceable(8) %0, %swift.refcounted* %1) #0 !dbg !836 {
  %3 = alloca %TSd, align 8
  call void @llvm.dbg.declare(metadata %TSd* %3, metadata !840, metadata !DIExpression()), !dbg !847
  %4 = bitcast %TSd* %3 to i8*, !dbg !848
  call void @llvm.lifetime.start.p0i8(i64 8, i8* %4), !dbg !848
  %5 = getelementptr inbounds %TSd, %TSd* %3, i32 0, i32 0, !dbg !848
  store double 0.000000e+00, double* %5, align 8, !dbg !848
  %6 = call swiftcc i8* @swift_autoDiffProjectTopLevelSubcontext(%swift.refcounted* %1) #8, !dbg !850
  %7 = bitcast i8* %6 to %"T6custom061_AD__$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSayK33G_tFySdzcfU__bb4__PB__src_0_wrt_033_E323A9EF754FCEC1F4377C8DDDEDCDFBLLV"*, !dbg !850
  %8 = getelementptr inbounds %"T6custom061_AD__$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSayK33G_tFySdzcfU__bb4__PB__src_0_wrt_033_E323A9EF754FCEC1F4377C8DDDEDCDFBLLV", %"T6custom061_AD__$s6custom19te
  %9 = bitcast %"T6custom061_AD__$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSayK35G_tFySdzcfU__bb4__Pred__src_0_wrt_033_E323A9EF754FCEC1F4377C8DDDEDCDFBLLO"* %8 to %"T6custom061_AD__$s6custom19testUpdat
  %10 = getelementptr inbounds %"T6custom061_AD__$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSayK33G_tFySdzcfU__bb2__PB__src_0_wrt_033_E323A9EF754FCEC1F4377C8DDDEDCDFBLLV", %"T6custom061_AD__$s6custom19t
  %11 = bitcast %"T6custom061_AD__$s6custom19testUpdateByCallingyyKF8fOfArrayL_5arraySdSayK35G_tFySdzcfU__bb2__Pred__src_0_wrt_033_E323A9EF754FCEC1F4377C8DDDEDCDFBLLO"* %10 to i8**, !dbg !850
  %12 = load i8*, i8** %11, align 8, !dbg !850
  %13 = getelementptr inbounds %TSd, %TSd* %0, i32 0, i32 0, !dbg !850
  %14 = load double, double* %13, align 8, !dbg !850
  call void @llvm.dbg.value(metadata double %14, metadata !841, metadata !DIExpression()), !dbg !851
...

So, !840 and !841 conflict here

[asl]

Ok, I think we're having some generic issue about debug info generation in differential swift code (and it looks similar to #55703 actually).

This is generated pullback:

// pullback of closure #1 in fOfArray #1 (array:) in testUpdateByCalling()
sil private [ossa] @$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr : $@convention(thin) (@inout Double, @guaranteed Builtin.NativeObject) -> () {
// %0                                             // users: %100, %23, %22
// %1                                             // user: %18
bb0(%0 : $*Double, %1 : @guaranteed $Builtin.NativeObject):
  %2 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %107, %100, %97, %56, %55, %5
  %3 = witness_method $Double, #AdditiveArithmetic.zero!getter : <Self where Self : AdditiveArithmetic> (Self.Type) -> () -> Self : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveAr
  %4 = metatype $@thick Double.Type, loc "custom.swift":52:47, scope 2 // user: %5
  %5 = apply %3<Double>(%2, %4) : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@thick τ_0_0.Type) -> @out τ_0_0, loc "custom.swift":52:47, scope 2
  %6 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %106, %105, %84, %83, %67, %54, %52, %51, %9
  %7 = witness_method $Double, #AdditiveArithmetic.zero!getter : <Self where Self : AdditiveArithmetic> (Self.Type) -> () -> Self : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveAr
  %8 = metatype $@thick Double.Type, loc "custom.swift":52:47, scope 2 // user: %9
  %9 = apply %7<Double>(%6, %8) : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@thick τ_0_0.Type) -> @out τ_0_0, loc "custom.swift":52:47, scope 2
  %10 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %104, %103, %84, %83, %56, %55, %54, %52, %51, %37, %36, %13
  %11 = witness_method $Double, #AdditiveArithmetic.zero!getter : <Self where Self : AdditiveArithmetic> (Self.Type) -> () -> Self : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveA
  %12 = metatype $@thick Double.Type, loc "custom.swift":52:47, scope 2 // user: %13
  %13 = apply %11<Double>(%10, %12) : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveArithmetic> (@thick τ_0_0.Type) -> @out τ_0_0, loc "custom.swift":52:47, scope 2
  %14 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %102, %101, %37, %36, %23, %22, %17
  %15 = witness_method $Double, #AdditiveArithmetic.zero!getter : <Self where Self : AdditiveArithmetic> (Self.Type) -> () -> Self : $@convention(witness_method: AdditiveArithmetic) <τ_0_0 where τ_0_0 : AdditiveA
  %16 = metatype $@thick Double.Type, loc "custom.swift":52:47, scope 2 // user: %17

Note that we're having multiple locations for element local variable that corresponds to argument number 1.

After SROA some of these alloca's got sliced:

  %2 = alloc_stack $Builtin.FPIEEE64, var, (name "element", loc "custom.swift":52:47, scope 2), argno 1, type $*Double, expr op_fragment:#Double._value, loc "<compiler-generated>":0:0, scope 2 // users: %110, %10
  %3 = integer_literal $Builtin.Int64, 0, loc "custom.swift":52:47, scope 2 // user: %4
  %4 = builtin "sitofp_Int64_FPIEEE64"(%3 : $Builtin.Int64) : $Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // user: %5
  %5 = struct $Double (%4 : $Builtin.FPIEEE64), loc "custom.swift":52:47, scope 2 // user: %6
  %6 = destructure_struct %5 : $Double, loc "custom.swift":52:47, scope 2 // user: %7
  store %6 to [trivial] %2 : $*Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // id: %7
  %8 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %90, %67, %63, %60, %93, %12, %109, %84
  %9 = integer_literal $Builtin.Int64, 0, loc "custom.swift":52:47, scope 2 // user: %10
  %10 = builtin "sitofp_Int64_FPIEEE64"(%9 : $Builtin.Int64) : $Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // user: %11
  %11 = struct $Double (%10 : $Builtin.FPIEEE64), loc "custom.swift":52:47, scope 2 // user: %12
  store %11 to [trivial] %8 : $*Double, loc "custom.swift":52:47, scope 2 // id: %12
  %13 = alloc_stack $Builtin.FPIEEE64, var, (name "element", loc "custom.swift":52:47, scope 2), argno 1, type $*Double, expr op_fragment:#Double._value, loc "<compiler-generated>":0:0, scope 2 // users: %108, %1
  %14 = integer_literal $Builtin.Int64, 0, loc "custom.swift":52:47, scope 2 // user: %15
  %15 = builtin "sitofp_Int64_FPIEEE64"(%14 : $Builtin.Int64) : $Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // user: %16
  %16 = struct $Double (%15 : $Builtin.FPIEEE64), loc "custom.swift":52:47, scope 2 // user: %17
  %17 = destructure_struct %16 : $Double, loc "custom.swift":52:47, scope 2 // user: %18
  store %17 to [trivial] %13 : $*Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // id: %18
  %19 = alloc_stack $Builtin.FPIEEE64, var, (name "element", loc "custom.swift":52:47, scope 2), argno 1, type $*Double, expr op_fragment:#Double._value, loc "<compiler-generated>":0:0, scope 2 // users: %107, %2
...

Note that sliced allocas got artificial location information. After some optimizations we're ending with:

 %2 = alloc_stack $Double, var, name "element", argno 1, loc "custom.swift":52:47, scope 2 // users: %40, %26, %6, %54, %36
  %3 = integer_literal $Builtin.Int64, 0, loc "custom.swift":52:47, scope 2 // user: %4
  %4 = builtin "sitofp_Int64_FPIEEE64"(%3 : $Builtin.Int64) : $Builtin.FPIEEE64, loc "custom.swift":52:47, scope 2 // users: %21, %16, %5
  %5 = struct $Double (%4 : $Builtin.FPIEEE64), loc "custom.swift":52:47, scope 2 // user: %6
  store %5 to %2 : $*Double, loc "custom.swift":52:47, scope 2 // id: %6
  %7 = builtin "autoDiffProjectTopLevelSubcontext"(%1 : $Builtin.NativeObject) : $Builtin.RawPointer, loc "custom.swift":52:44, scope 2 // user: %8
  %8 = pointer_to_address %7 : $Builtin.RawPointer to [strict] $*_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4__PB__src_0_wrt_0, loc "custom.swift":52:44, scope 2 // user: %9
  %9 = struct_element_addr %8 : $*_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4__PB__src_0_wrt_0, #_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4__PB
  %10 = load %9 : $*_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4__Pred__src_0_wrt_0, loc "custom.swift":52:44, scope 2 // user: %13
  %11 = struct_element_addr %0 : $*Double, #Double._value, loc "custom.swift":52:44, scope 2 // user: %12
  %12 = load %11 : $*Builtin.FPIEEE64, loc "custom.swift":52:44, scope 2 // users: %15, %21
  %13 = unchecked_enum_data %10 : $_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4__Pred__src_0_wrt_0, #_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb4_
  %14 = struct_extract %13 : $_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb2__PB__src_0_wrt_0, #_AD__$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU__bb2__PB__sr
  debug_value %12 : $Builtin.FPIEEE64, var, (name "element", loc "custom.swift":52:47, scope 2), argno 1, type $*Double, expr op_fragment:#Double._value, loc "<compiler-generated>":0:0, scope 2 // id: %15
...

So, what happens is that we're having different debug info for argument number 1: alloc_stack with explicit location information and debug_value with artificial. The thing is: these are separate locations that got tied to the same argument. Actually, the variable name is misleading here. These are different adjoint buffers for the same value in different basic blocks. Therefore I think:

1. We need to drop "argno" here entirely
2. We need to change variable name to indicate that these are adjoints

@BradLarson @dan-zheng @rxwei Thoughts?