shrink-wrap: fix shrink-wrapping for no-return paths

[thegameg]

Summary

Shrink-wrapping uses post-dominators to find a restore point that post-dominates all the uses of CSR / stack.

The way dominator trees are modeled in LLVM today is that unreachable blocks are not present in a generic dominator tree. I am actually guessing this from some comments: include/llvm/Support/GenericDomTree.h:423,
So, an unreachable node is dominated by anything: include/llvm/Support/GenericDomTree.h:467.
There seems to be an explanation that an unreachable block may or may not appear in the dominator tree: include/llvm/IR/Dominators.h:103:

/// A forward unreachable block may appear in the dominator tree, or it may
/// not.  If it does, dominance queries will return results as if all reachable
/// blocks dominate it.  When asking for a Node corresponding to a potentially
/// unreachable block, calling code must handle the case where the block was
/// unreachable and the result of getNode() is nullptr.
///
/// Generally, a block known to be unreachable when the dominator tree is
/// constructed will not be in the tree.  One which becomes unreachable after
/// the dominator tree is initially constructed may still exist in the tree,
/// even if the tree is properly updated. Calling code should not rely on the
/// preceding statements; this is stated only to assist human understanding.

From that, I assume the MachineDominators might work in the same way.

Since for post-dominators, a no-return block is considered "unreachable", calling findNearestCommonDominator on an unreachable node A and a non-unreachable node B, will return B, which can be false.

I suggest that if we find such a node, we bail out since there is no good restore point.

[qcolombet]

LGTM

[MatzeB]

Good find Francis. I am not convinced however the fix is save (see below). And for that matter the existing if (!Save) code path in FindIDom() seems to suffer the same problem.

lib/CodeGen/ShrinkWrap.cpp

292

Setting this to nullptr seems incorrect to me. From my understanding that resets us to a state where it seems like we haven't looked at any part of the function and will just pick the next block that comes along, which seems really sketchy to me. Shouldn't we rather abort the whole shrinkwrapping process at this point?

[MatzeB]

Okay I misread the code, looks like ArePointsInteresting() abort if Save or Restore is nullptr. So this seems fine.

Would be nice to add a comment that setting it to nullptr will abort the shrinkwrapping when committing.

[thegameg]

Sorry, I forgot to add the review link to the commit message. I added comments, fixed a Thumb2 test (test/CodeGen/Thumb2/v8_IT_5.ll) and committed as r303130.

lib/CodeGen/ShrinkWrap.cpp

void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB,
 
   if (!Save) {
     DEBUG(dbgs() << "Found a block that is not reachable from Entry\n");
     return;
   }
 
   if (!Restore)
     Restore = &MBB;
-  else
+  else if (MPDT->getNode(&MBB)) // If the block is not in the post dom tree, it
+                                // means the block never returns. If that's the
+                                // case, we don't want to call
+                                // `findNearestCommonDominator`, which will
+                                // return `Restore`.
     Restore = MPDT->findNearestCommonDominator(Restore, &MBB);
+  else
+    Restore = nullptr;
 
   // Make sure we would be able to insert the restore code before the
   // terminator.
   if (Restore == &MBB) {
     for (const MachineInstr &Terminator : MBB.terminators()) {
       if (!useOrDefCSROrFI(Terminator, RS))
         continue;
       // One of the terminator needs to happen before the restore point.

test/CodeGen/X86/x86-shrink-wrapping.ll

if.else:                                          ; preds = %entry
   %mul = shl nsw i32 %N, 1
   br label %if.end
 
 if.end:                                           ; preds = %if.else, %for.end
   %sum.1 = phi i32 [ %shl, %for.end ], [ %mul, %if.else ]
   ret i32 %sum.1
 }
 
-declare void @somethingElse(...)
-
 ; Check with a more complex case that we do not have restore within the loop and
 ; save outside.
 ; CHECK-LABEL: loopInfoRestoreOutsideLoop:
 ;
 ; ENABLE: testl %edi, %edi
 ; ENABLE-NEXT: je [[ELSE_LABEL:LBB[0-9_]+]]
 ;
 ; Prologue code.

 
 for.inc:
   %inc18.i = add nuw nsw i32 %p13, 1
   %cmp = icmp slt i32 %inc18.i, 7
   br i1 %cmp, label %for.cond8, label %fn1.exit
 }
 
 attributes #4 = { "no-frame-pointer-elim"="true" }
+
+@x = external global i32, align 4
+@y = external global i32, align 4
+
+; The post-dominator tree does not include the branch containing the infinite
+; loop, which can occur into a misplacement of the restore block, if we're
+; looking for the nearest common post-dominator of an "unreachable" block.
+
+; CHECK-LABEL: infiniteLoopNoSuccessor:
+; CHECK: ## BB#0:
+; Make sure the prologue happens in the entry block.
+; CHECK-NEXT: pushq %rbp
+; ...
+; Make sure we don't shrink-wrap.
+; CHECK: ## BB#1
+; CHECK-NOT: pushq %rbp
+; ...
+; Make sure the epilogue happens in the exit block.
+; CHECK: ## BB#5
+; CHECK: popq %rbp
+; CHECK-NEXT: retq
+define void @infiniteLoopNoSuccessor() #5 {
+  %1 = load i32, i32* @x, align 4
+  %2 = icmp ne i32 %1, 0
+  br i1 %2, label %3, label %4
+
+; <label>:3:
+  store i32 0, i32* @x, align 4
+  br label %4
+
+; <label>:4:
+  call void (...) @somethingElse()
+  %5 = load i32, i32* @y, align 4
+  %6 = icmp ne i32 %5, 0
+  br i1 %6, label %10, label %7
+
+; <label>:7:
+  %8 = call i32 (...) @something()
+  br label %9
+
+; <label>:9:
+  call void (...) @somethingElse()
+  br label %9
+
+; <label>:10:
+  ret void
+}
+
+declare void @somethingElse(...)
+
+attributes #5 = { nounwind  "no-frame-pointer-elim-non-leaf" }