SILCombine for differential_function creates SIL with invalid ownership when enabling OSSA modules

[eeckstein]

Description

When running the test AutoDiff/validation-test/reabstraction.swift in optimize mode and with sil-verify-all, the compiler crashes with

Begin Error in Function: '$s4mainyycfU7_'
Found over consume?!
Value:   %9 = partial_apply [callee_guaranteed] %8(%3) : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <Float>) // users: %20, %14
User:   %14 = convert_function %9 : $@callee_guaranteed (@in_guaranteed Float) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <Float>) to $@callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <τ_0_1>) for <Float, Float> // user: %16
Block: bb0
Consuming Users:
  %20 = convert_function %9 : $@callee_guaranteed (@in_guaranteed Float) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <Float>) to $@callee_guaranteed (@in_guaranteed Float) -> (Float, @owned @callee_guaranteed (@in_guaranteed Float) -> Float) // user: %22
  %14 = convert_function %9 : $@callee_guaranteed (@in_guaranteed Float) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <Float>) to $@callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@in_guaranteed τ_0_0) -> Float for <τ_0_1>) for <Float, Float> // user: %16

Reproduction

• Enable OSSA modules, e.g. by cherry-picking cb4db5a
• Enable sil-verify-all, e.g. by

diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
index 2a5a5c7d09e..7e1541556c9 100644
--- a/lib/Frontend/CompilerInvocation.cpp
+++ b/lib/Frontend/CompilerInvocation.cpp
@@ -2881,7 +2881,7 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
                    OPT_disable_lifetime_dependence_diagnostics,
                    Opts.EnableLifetimeDependenceDiagnostics);
 
-  Opts.VerifyAll |= Args.hasArg(OPT_sil_verify_all);
+  Opts.VerifyAll = true; // |= Args.hasArg(OPT_sil_verify_all);
   Opts.VerifyNone |= Args.hasArg(OPT_sil_verify_none);
   Opts.VerifyOwnershipAll |= Args.hasArg(OPT_sil_ownership_verify_all);
   Opts.DebugSerialization |= Args.hasArg(OPT_sil_debug_serialization);

• Run the test AutoDiff/validation-test/reabstraction.swift in optimized mode, e.g. with ninja check-swift-optimize

Expected behavior

no crash

Environment

Additional information

No response

[eeckstein]

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

[eeckstein]

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

[asl]

Testcase is:

  typealias FnTy<T: Differentiable> = @differentiable(reverse) (T) -> Float
  func id<T: Differentiable>(_ f: @escaping FnTy<T>) -> FnTy<T> { f }
  func inner(_ x: Float) -> Float {
    7 * x * x
  }
  let transformed = id(inner)
  let _ = gradient(at: 3, of: transformed)

[asl]

And the problem turned out to be pretty trivial. At some point after some inlining we're ending with the following code:

  %15 = differentiable_function [parameters 0] [results 0] %8 with_derivative {%11, %14} // users: %16, %20
  %16 = begin_borrow [lexical] %15                // users: %18, %17, %19
  debug_value %16, let, name "f", argno 1         // id: %17
  %18 = copy_value %16                            // user: %21
  end_borrow %16                                  // id: %19
  destroy_value %15                               // id: %20
  %21 = convert_function %18 to $@differentiable(reverse) @callee_guaranteed (@in_guaranteed Float) -> Float // users: %36, %22

So, differentiable_function has two uses one being purely debug one, but hidden behind borrow brackets. It seems that for this particular peephole there was no # of users check. The peephole pushed conversion before differentiable_function, however, as we're essentially replacing convert_function with newly built differentiable_function, the old one remained and after couple of more peepholes screwed the things around.

So, the solution is to add the check if we're having a single non-debug user. The borrow with debug_value inside is removed somewhere else and the peephole is fired just a little bit later.

Also, it seems that some conversion peepholes around use strong hasSingleUse() and any debug value usage prevents them from being used. I decided to fix it while there.