The AutodiffClosureSpecialization pass creates SIL with invalid ownership when enabling OSSA modules

[eeckstein]

Description

When running the test AutoDiff/validation-test/array.swift in optimize mode, the compiler crashes with

Begin Error in Function: '$sSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVyADySf_G_GA2FIeggoo_TJSpSSUpSrUSUP07$sSa16_a4AA14b47RzlE15_vjpConcatenateySayxG5value_SaA2aBRzlE0B4c54Vy13TangentVectorQz_G_AJtAJc8pullbacktAD_ADtFZAKL_yAJ_v3AjaO11FSaySfG_Tg5SaySaySfGGAITf1nc_n'
Found over consume?!
Value:   %4 = partial_apply [callee_guaranteed] %3(%1, %2) : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed Array<Array<Float>>, @guaranteed Array<Array<Float>>) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // users: %11, %6
User:   %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed @callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView)) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // users: %12, %7
Block: bb0
Consuming Users:
  destroy_value %4 : $@callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // id: %11
  %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed @callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView)) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // users: %12, %7

Reproduction

• Enable OSSA modules, e.g. by cherry-picking cb4db5a
• Run the test AutoDiff/validation-test/array.swift in optimized mode, e.g. with ninja check-swift-optimize

Expected behavior

no crash

Environment

Additional information

No response

[eeckstein]

@jkshtj would you be able to take care of this?

[eeckstein]

I'm temporarily disabling this optimization for OSSA: #78849

[jkshtj]

@jkshtj would you be able to take care of this?

@eeckstein I haven't been working on Swift Autodiff for a while now. May be @asl can?

[asl]

Tagging @kovdan01 @JaapWijnen

[asl]

Reduced testcase:

import _Differentiation

typealias FloatArrayTan = Array<Float>.TangentVector

func controlFlowNestedLiteral(
  _ x: Float, _ y: Float, _ bool: Bool = true
) -> [Float] {
  var result: [[Float]] = []
  var result2 = bool ? result + [[x]] : result + [[x]]
  var result3 = bool ? (bool ? result2 + [[y]] : result2 + [[y]]) : result2 + [[y]]
  return result3[0] + [result3[1][0]]
}
let pb = pullback(at: 3, 4, of: { controlFlowNestedLiteral($0, $1) })

[asl]

The specialized parameter thunk generated is clearly incorrect:

// specialized autodiff subset parameters thunk for pullback from @escaping @callee_guaranteed (@guaranteed [[Float]<A>.DifferentiableView]<A>.DifferentiableView) -> (@owned [[Float]<A>.DifferentiableView]<A>.DifferentiableView, @owned [[Float]<A>.DifferentiableView]<A>.DifferentiableView)
sil shared [transparent] [thunk] [ossa] @$sSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVyADySf_G_GA2FIeggoo_TJSpSSUpSrUSUP07$sSa16_a4AA14b47RzlE15_vjpConcatenateySayxG5value_SaA2aBRzlE0B4c54Vy13TangentVectorQz_G_AJtAJc8pullbacktAD_ADtFZAKL_yAJ_v3AjaO11FSaySfG_Tg5SaySaySfGGAITf1nc_n : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>>, @owned Array<Array<Float>>) -> @owned Array<Array<Float>.DifferentiableView>.DifferentiableView {
// %0                                             // user: %7
// %1                                             // user: %4
// %2                                             // user: %4
bb0(%0 : @guaranteed $Array<Array<Float>.DifferentiableView>.DifferentiableView, %1 : @owned $Array<Array<Float>>, %2 : @owned $Array<Array<Float>>):
  // function_ref specialized pullback #1 <A>(_:) in static Array<A>._vjpConcatenate(_:_:)
  %3 = function_ref @$sSa16_DifferentiationAA14DifferentiableRzlE15_vjpConcatenateySayxG5value_SaA2aBRzlE0B4ViewVy13TangentVectorQz_G_AJtAJc8pullbacktAD_ADtFZAKL_yAJ_AJtAjaBRzlFSaySfG_Tg5 : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed Array<Array<Float>>, @guaranteed Array<Array<Float>>) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // user: %4
  %4 = partial_apply [callee_guaranteed] %3(%1, %2) : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed Array<Array<Float>>, @guaranteed Array<Array<Float>>) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // users: %11, %6
  // function_ref specialized thunk for @escaping @callee_guaranteed (@guaranteed [A.Differentiable.TangentVector]<A>.DifferentiableView) -> (@owned [A.Differentiable.TangentVector]<A>.DifferentiableView, @owned [A.Differentiable.TangentVector]<A>.DifferentiableView)
  %5 = function_ref @$sSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVy13TangentVectorAaBPQz_GA2HIeggoo_A3HIeggoo_AaBRzlTRSaySfG_Tg5 : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed @callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView)) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // user: %6
  %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView, @guaranteed @callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView)) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // users: %12, %7
  %7 = apply %6(%0) : $@callee_guaranteed (@guaranteed Array<Array<Float>.DifferentiableView>.DifferentiableView) -> (@owned Array<Array<Float>.DifferentiableView>.DifferentiableView, @owned Array<Array<Float>.DifferentiableView>.DifferentiableView) // user: %8
  (%8, %9) = destructure_tuple %7                 // users: %10, %13
  destroy_value %8                                // id: %10
  destroy_value %4                                // id: %11
  destroy_value %6                                // id: %12
  return %9                                       // id: %13
} // end sil function '$sSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVyADySf_G_GA2FIeggoo_TJSpSSUpSrUSUP07$sSa16_a4AA14b47RzlE15_vjpConcatenateySayxG5value_SaA2aBRzlE0B4c54Vy13TangentVectorQz_G_AJtAJc8pullbacktAD_ADtFZAKL_yAJ_v3AjaO11FSaySfG_Tg5SaySaySfGGAITf1nc_n'

So, here %4 is having two consuming uses.

[asl]

So, the ownership-tracking logic at
https://github.com/swiftlang/swift/blob/main/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift#L788
and
https://github.com/swiftlang/swift/blob/main/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift#L981

is broken.

Essentially, given the sequence of partial_apply instructions like in #78847 (comment) it wants to explicitly release all of them. This is certainly not correct as in example above %6 consumes %4. So we only need to release the final closure that we're applying here, and not the intermediate ones as these are consumed during reabstraction conversion.