[llvm] [MachineSink] Fix missing sinks along critical edges (PR #97618)

[Min-Yih Hsu via llvm-commits]

mshockwave wrote:

Hi @alexfh thanks for discovering this bug. I think you just found a really interesting issue caused by an old bug.

Here is the MachineFunction right before the crash:
```
# Machine code for function _f: IsSSA, TracksLiveness
Function Live Ins: $x0 in %6, $w1 in %7

bb.0.entry:
  successors: %bb.1(0x30000000), %bb.2(0x50000000); %bb.1(37.50%), %bb.2(62.50%)
  liveins: $x0, $w1
  %7:gpr32 = COPY $w1
  %6:gpr64common = COPY $x0
  %8:gpr32 = COPY %7:gpr32
  %9:gpr32 = LDRSBWui %6:gpr64common, 0 :: (load (s8) from %ir.project)
  %11:gpr64all = IMPLICIT_DEF
  %10:gpr64 = INSERT_SUBREG %11:gpr64all(tied-def 0), %9:gpr32, %subreg.sub_32
  %12:gpr64sp = ANDXri killed %10:gpr64, 4103
  %0:gpr64 = COPY %12:gpr64sp
  %1:gpr64all = COPY %12:gpr64sp
  TBZW %9:gpr32, 31, %bb.2
  B %bb.1

bb.1.cond.true.i9.i.i:
; predecessors: %bb.0
  successors: %bb.5(0x40000000), %bb.4(0x40000000); %bb.5(50.00%), %bb.4(50.00%)

  %13:gpr64 = LDRXui %6:gpr64common, 0 :: (load (s64) from %ir.project)
  %2:gpr64 = COPY %13:gpr64
  TBNZW %8:gpr32, 0, %bb.5
  B %bb.4

bb.2._ZNKSt3__u12basic_stringIcNS_11char_traitsIcEENS_9allocatorIcEEE16__get_short_sizeEv.exit.i.i.i:
; predecessors: %bb.0
  successors: %bb.5(0x40000000), %bb.4(0x40000000); %bb.5(50.00%), %bb.4(50.00%)

  %18:gpr64 = COPY %1:gpr64all
  TBNZW %8:gpr32, 0, %bb.5
  B %bb.4

bb.3.for.cond.i:
; predecessors: %bb.4
  successors: %bb.5(0x04000000), %bb.4(0x7c000000); %bb.5(3.12%), %bb.4(96.88%)

  %15:gpr64sp = ADDXri %5:gpr64sp, 1, 0
  %4:gpr64all = COPY %15:gpr64sp
  CBZX %19:gpr64, %bb.5
  B %bb.4

bb.4.for.body.i:
; predecessors: %bb.1, %bb.2, %bb.3
  successors: %bb.6(0x04000000), %bb.3(0x7c000000); %bb.6(3.12%), %bb.3(96.88%)

  %19:gpr64 = PHI %19:gpr64, %bb.3, %2:gpr64, %bb.1, %18:gpr64, %bb.2
  %5:gpr64sp = PHI %6:gpr64common, %bb.1, %4:gpr64all, %bb.3, %6:gpr64common, %bb.2
  %14:gpr32 = LDRBBui %5:gpr64sp, 0 :: (load (s8) from %ir.__first.sroa.0.010.i)
  CBZW killed %14:gpr32, %bb.6
  B %bb.3

bb.5.common.ret:
; predecessors: %bb.1, %bb.2, %bb.3, %bb.6, %bb.7

  RET_ReallyLR

bb.6.cond.false:
; predecessors: %bb.4
  successors: %bb.5(0x30000000), %bb.7(0x50000000); %bb.5(37.50%), %bb.7(62.50%)

  %16:gpr32 = COPY %1.sub_32:gpr64all
  TBNZW killed %16:gpr32, 7, %bb.5
  B %bb.7

bb.7.cond.end.i.thread.i.i72:
; predecessors: %bb.6
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  STRXui %0:gpr64, %6:gpr64common, 0 :: (store (s64) into %ir.project)
  B %bb.5

# End machine code for function _f.
```

MachineSink thinks it's profitable to sink `%0:gpr64 = COPY %12:gpr64sp` from bb.0 to bb.4 after the critical edge between those two blocks are broken. But bb.4 is not even a successor of bb.0, hence the assertion failure.

The reason it thought bb.0 - bb.4 is even a critical edge in the first place is because bb.4 is the `SuccToSinkTo` block w.r.t the current block (i.e. bb.0) returned from `FindSuccToSinkTo`. MachineSink consider the edge between current block and `SuccToSinkTo` a critical edge if `SuccToSinkTo` has more than one predecessor -- which is true for most of the cases. However, `FindSuccToSinkTo` not only return a (immediate) successor block but also look beyond to all descendant blocks: https://github.com/llvm/llvm-project/blob/43024a465120315464490ca9c6074553cd968089/llvm/lib/CodeGen/MachineSink.cpp#L1202  . Which is why `FindSuccToSinkTo` return bb.4 despite the fact that bb.4 is not even a successor of bb.0. (The reason I said this is an old bug because the logics of looking into descendant blocks has been there for a long time)

I think a simple fix on `isLegalToBreakCriticalEdge` to reject cases like this should fix the issue. I'm cooking a patch right now...

https://github.com/llvm/llvm-project/pull/97618