[ARM] Only update the successor edges for immediate predecessors of PrologueMBB.

[ZhiyaoMa98]

Summary

When adjusting the function prologue for segmented stacks, only update the successor edges of the immediate predecessors of the original prologue.

I encountered a segmentation fault in LLVM while compiling the Rust core targeting armv7em with segmented stack enabled. Below is my attempt to fix it. I am not very certain about the correctness of my patch, because I do not fully understand the code.

Looking at the do-while loop

SmallPtrSet<MachineBasicBlock *, 8> BeforePrologueRegion;
SmallVector<MachineBasicBlock *, 2> WalkList;
WalkList.push_back(&PrologueMBB);

do {
  MachineBasicBlock *CurMBB = WalkList.pop_back_val();
  for (MachineBasicBlock *PredBB : CurMBB->predecessors()) {
    if (BeforePrologueRegion.insert(PredBB).second)
      WalkList.push_back(PredBB);
  }
} while (!WalkList.empty());

it seems BeforePrologueRegion contains non-immediate predecessors of PrologueMBB .

However, in the following while loop, ReplaceUsesOfBlockWith assumes that MBB is an immediate predecessor of PrologueMBB.

for (MachineBasicBlock *MBB : BeforePrologueRegion) {
  // Make sure the LiveIns are still sorted and unique.
  MBB->sortUniqueLiveIns();
  // Replace the edges to PrologueMBB by edges to the sequences
  // we are about to add.
  MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);
}

The assumption comes from the function call chain

MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);
--> replaceSuccessor(Old, New);
----> Old->removePredecessor(this);
------> assert(I != Predecessors.end()); Predecessors.erase(I);

I ran into this problem when I was compiling Rust core in release mode targeting armv7em with segmented stack enabled. The Predecessors.erase(I); line led to a segmentation fault.

I fixed the problem by adding the if-continue statement. It still passed all of the tests.

[john.brawn]

I don't know much about this specific bit of code, but a comment in ReplaceUsesOfBlockWith mentions that it's expected that 'Old' is a successor of 'this' so the fix looks reasonable. However, could you also add a test for this?

llvm/lib/Target/ARM/ARMFrameLowering.cpp

2564–2569

I think this would be a bit clearer if rewritten as

if (MBB->isSuccessor(&PrologueMBB))
  MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);

i.e. the test and action are both MBB->thing(&PrologueMBB).

[ZhiyaoMa98]

I would like to add a test, but I do not know how. My understanding is that I will need to create an example that PrologueMBB has non-immediate predecessors. Is there any example I can follow? Thanks!

[john.brawn]

In D122959#3475470, @ZhiyaoMa98 wrote:

I would like to add a test, but I do not know how. My understanding is that I will need to create an example that PrologueMBB has non-immediate predecessors. Is there any example I can follow? Thanks!

It may be easiest to start with the original file they you saw the segfault. Adding "--emit llvm-ir" to the rustc command-line will cause a .ll file to be produced, then doing "bugpoint -run-llc file.ll" should hopefully reduce that down to a simple test case.

[ZhiyaoMa98]

I successfully reproduced the segfault by llc -relocation-model=ropi-rwpi large_llvm_ir.ll
I need to pass the -relocation-model=ropi-rwpi flag to bugpoint, any hint on this?

Edit: it seems to be --tool-args, trying this now.

[ZhiyaoMa98]

Added tests. Before fix it segfaults on the new tests. After fix it just runs through.

[ZhiyaoMa98]

Addressed inline comment to make the code more readable.

[ZhiyaoMa98]

I don't have commit access. Please credit me as follows if everything looks good. Thanks.
Name: Zhiyao Ma
Email: [e-mail]

[john.brawn]

Looks good, I'll commit this for you.

llvm/lib/Target/ARM/ARMFrameLowering.cpp

void ARMFrameLowering::adjustForSegmentedStacks(
   for (MachineBasicBlock *B : AddedBlocks) {
     BeforePrologueRegion.erase(B);
     MF.insert(PrologueMBB.getIterator(), B);
   }
 
   for (MachineBasicBlock *MBB : BeforePrologueRegion) {
     // Make sure the LiveIns are still sorted and unique.
     MBB->sortUniqueLiveIns();
     // Replace the edges to PrologueMBB by edges to the sequences
-    // we are about to add.
+    // we are about to add, but only update for immediate predecessors.
+    if (MBB->isSuccessor(&PrologueMBB))
       MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);
   }
 
   // The required stack size that is aligned to ARM constant criterion.
   AlignedStackSize = alignToARMConstant(StackSize);
 
   // When the frame size is less than 256 we just compare the stack
   // boundary directly to the value of the stack pointer, per gcc.
   bool CompareStackPointer = AlignedStackSize < kSplitStackAvailable;
 
   // We will use two of the callee save registers as scratch registers so we

llvm/test/CodeGen/ARM/segmented-stacks.ll

 ; ARM-linux-LABEL: test_sibling_call_empty_frame:
 ; ARM-linux:      bl      __morestack
 
 ; ARM-android-LABEL: test_sibling_call_empty_frame:
 ; ARM-android:      bl      __morestack
 
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; ARM-linux-LABEL: build_should_not_segfault:
+; ARM-android-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }

llvm/test/CodeGen/Thumb/segmented-stacks.ll

 
 ; Thumb-android-LABEL: test_nostack:
 ; Thumb-android-NOT:   bl __morestack
 
 ; Thumb-linux-LABEL: test_nostack:
 ; Thumb-linux-NOT:   bl __morestack
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; Thumb-linux-LABEL: build_should_not_segfault:
+; Thumb-android-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }

llvm/test/CodeGen/Thumb2/segmented-stacks.ll

 ; ARM-NEXT:    ldm     sp!, {lr}
 ; ARM-NEXT:    pop     {r4, r5}
 ; ARM-NEXT:    bx      lr
 
 ; ARM:         .LCPI2_0:
 ; ARM-NEXT:    .long   40192
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; THUMB-LABEL: build_should_not_segfault:
+; ARM-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }