[SCCP] do not clean up dead blocks that have their address taken

[nickdesaulniers]

Summary

[SCCP] do not clean up dead blocks that have their address taken

Fixes a crash observed in IPSCCP.

Because the SCCPSolver has already internalized BlockAddresses as
Constants or ConstantExprs, we don't want to try to update their Values
in the ValueLatticeElement. Instead, continue to propagate these
BlockAddress Constants, continue converting BasicBlocks to unreachable,
but don't delete the "dead" BasicBlocks which happen to have their
address taken. Leave replacing the BlockAddresses to another pass.

Fixes: https://github.com/llvm/llvm-project/issues/54238
Fixes: https://github.com/llvm/llvm-project/issues/54251

[nickdesaulniers]

llvm/lib/Transforms/Scalar/SCCP.cpp

527–530

There's probably some code somewhere in llvm::SimplifyInstruction that already does this transform (as observed by SimplifyCFG in https://github.com/llvm/llvm-project/issues/54328#issuecomment-1067229660). maybe that can be reused here?

Since the behavior is surprising to users (https://github.com/llvm/llvm-project/issues/54328) maybe we could consolidate such transforms with a comment about why this is done, for historical context?

llvm/lib/Transforms/Utils/SCCPSolver.cpp

424–426 ↗

(On Diff #415599)

This added bitcast is a bit curious, I'd like to clean it up, but I'm not sure yet how best to do so.

Basically, we know we want to replace i8* blockaddress(@fn, %bb) with i8* inttoptr (i32 1 to i8*), but the User of the BA and entry in ValueState is %1 = bitcast i8* blockaddress(@fn, %bb) to i64*.

Maybe I should be doing an operand replacement? As in take Old and replace its operand with New?

[nickdesaulniers]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

424–426 ↗

(On Diff #415599)

Looks like User::replaceUsesOfWith can't be called with a Constant. Maybe I need to construct a new Instruction to replace the prior use, with the new operand?

[nickdesaulniers]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

424–426 ↗

(On Diff #415599)

Looks like I need to be doing checks for ConstantExprs in the caller rather than dealing strictly with Constants.

[nickdesaulniers]

• handle ConstExpr's rather than just BitCasts

[nickdesaulniers]

llvm/lib/Transforms/Scalar/SCCP.cpp

533–542

While this removes the special case for bitcasts, it's still subtly not quite right; it only checks for one level of nesting of ConstExprs; there could be multiple "levels of nesting" for the aggregated Constant stored in the lattice.

i.e. the BlockAddress of interest might have the following User: i64* bitcast (i8* blockaddress(@main, %redirected) to i64*); but I worry if the lattice had something like i8* bitcast (i64* bitcast (i8* blockaddress(@main, %redirected) to i64*) to i8*); then we'd be trying to replace the former and not finding the latter. Unless the latter is also considered a User of the BlockAddress?

I can test that, but I'm going to see first whether we can query the solver whether it has reach-ability info ready BEFORE we internalize new ConstantExpr's containing BlockAddress Constants into the lattice. Then we could avoid this find+replace dance and simply do less work.

[nickdesaulniers]

llvm/lib/Transforms/Scalar/SCCP.cpp

533–542

but I'm going to see first whether we can query the solver whether it has reach-ability info ready BEFORE we internalize new ConstantExpr's containing BlockAddress Constants into the lattice. Then we could avoid this find+replace dance and simply do less work.

Yeah, the SCCPVisitor marks blocks executable during a traversal, while eagerly adding values to the lattice; so block executability is not precise until we've completed the traversal. At that point, it's too late as we've already merged values into the lattice.

[efriedma]

I think my preferred solution here would be to delay deleting dead blocks further, until the lattice stops mattering.

Any solution that involves trying to explicitly fix up the lattice, like the current patch, is going to cause problems for correctness and perf, I think.

For your other proposed solution, I don't think you can tell whether a block is reachable at the point where the blockaddress shows up in the lattice. You could blacklist all blockaddresses from the lattice, I guess, for a minor loss of optimization power.

You could also possibly use a TrackingVH in ValueLatticeElement; not sure what consequences that would have.

[nickdesaulniers]

In D121744#3386887, @efriedma wrote:

I think my preferred solution here would be to delay deleting dead blocks further, until the lattice stops mattering.

diff --git a/llvm/lib/Transforms/Scalar/SCCP.cpp b/llvm/lib/Transforms/Scalar/SCCP.cpp
index 794d4c5b5bc6..e9e3e90df875 100644
--- a/llvm/lib/Transforms/Scalar/SCCP.cpp
+++ b/llvm/lib/Transforms/Scalar/SCCP.cpp
@@ -505,7 +505,7 @@ bool llvm::runIPSCCP(
 
         MadeChanges = true;
 
-        if (&BB != &F.front())
+        if (&BB != &F.front() && !BB.hasAddressTaken())
           BlocksToErase.push_back(&BB);
         continue;
       }

?

[efriedma]

llvm/lib/Transforms/Scalar/SCCP.cpp

533–534

Something more like this, I think? Assuming this doesn't have any side-effects I'm not thinking of.

[nickdesaulniers]

llvm/lib/Transforms/Scalar/SCCP.cpp

533–534

yeah, that's the simplest fix.

I'm going to pre-commit updating llvm/test/Transforms/SCCP/dangling-block-address.ll's RUN line to use NPM, and run it through update_test_checks.py, just so we can see what changes.

[nickdesaulniers]

• rebase on pre-committed test changes. Use @efriedma's sugguestion.

[nikic]

LGTM. We could still delay these blocks later (after all functions have been handled), but I don't think it would be worthwhile.

llvm/lib/Transforms/Scalar/SCCP.cpp

for (Function &F : M) {
     // nodes in executable blocks we found values for. The function's entry
     // block is not part of BlocksToErase, so we have to handle it separately.
     for (BasicBlock *BB : BlocksToErase) {
       NumInstRemoved += changeToUnreachable(BB->getFirstNonPHI(),
                                             /*PreserveLCSSA=*/false, &DTU);
     }
     if (!Solver.isBlockExecutable(&F.front()))
       NumInstRemoved += changeToUnreachable(F.front().getFirstNonPHI(),
                                             /*PreserveLCSSA=*/false, &DTU);
 
     for (BasicBlock &BB : F)
       MadeChanges |= removeNonFeasibleEdges(Solver, &BB, DTU);
 
     for (BasicBlock *DeadBB : BlocksToErase)
+      if (!DeadBB->hasAddressTaken())
         DTU.deleteBB(DeadBB);
 
     for (BasicBlock &BB : F) {
       for (Instruction &Inst : llvm::make_early_inc_range(BB)) {
         if (Solver.getPredicateInfoFor(&Inst)) {
           if (auto *II = dyn_cast<IntrinsicInst>(&Inst)) {
             if (II->getIntrinsicID() == Intrinsic::ssa_copy) {
               Value *Op = II->getOperand(0);
               Inst.replaceAllUsesWith(Op);
               Inst.eraseFromParent();
             }
           }
         }
       }
     }
   }
 

llvm/test/Transforms/SCCP/dangling-block-address.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=internalize,ipsccp %s | FileCheck %s
 ; PR5569
 
-; IPSCCP should prove that the blocks are dead and delete them, and
-; properly handle the dangling blockaddress constants.
+; If BasicBlocks are unreachable, IPSCCP should convert them to unreachable.
+; Unless they have their address taken, delete them. If they do have their
+; address taken, let other passes replace these "dead" blockaddresses with some
+; other Constant.
+
+; CHECK: @bar.l = internal constant [2 x i8*] [i8* blockaddress(@bar, %lab0), i8* blockaddress(@bar, %end)]
 
-; CHECK: @bar.l = internal constant [2 x i8*] [i8* inttoptr (i32 1 to i8*), i8* inttoptr (i32 1 to i8*)]
 
 @code = global [5 x i32] [i32 0, i32 0, i32 0, i32 0, i32 1], align 4 ; <[5 x i32]*> [#uses=0]
 @bar.l = internal constant [2 x i8*] [i8* blockaddress(@bar, %lab0), i8* blockaddress(@bar, %end)] ; <[2 x i8*]*> [#uses=1]
 
 define void @foo(i32 %x) nounwind readnone {
 ; CHECK-LABEL: @foo(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    unreachable
 ;
 entry:
   %b = alloca i32, align 4                        ; <i32*> [#uses=1]
   store volatile i32 -1, i32* %b
   ret void
 }
 
 define void @bar(i32* nocapture %pc) nounwind readonly {
 ; CHECK-LABEL: @bar(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    unreachable
+; CHECK:       lab0:
+; CHECK-NEXT:    unreachable
+; CHECK:       end:
+; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %indirectgoto
 
 lab0:                                             ; preds = %indirectgoto
   %indvar.next = add i32 %indvar, 1               ; <i32> [#uses=1]
   br label %indirectgoto
 

 define i32 @main() nounwind readnone {
 ; CHECK-LABEL: @main(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    unreachable
 ;
 entry:
   ret i32 0
 }
+
+; https://github.com/llvm/llvm-project/issues/54238
+; https://github.com/llvm/llvm-project/issues/54251
+; https://github.com/llvm/llvm-project/issues/54328
+define i32 @test1() {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:    unreachable
+; CHECK:       redirected:
+; CHECK-NEXT:    unreachable
+;
+  %1 = bitcast i8* blockaddress(@test1, %redirected) to i64*
+  call void @set_return_addr(i64* %1)
+  ret i32 0
+
+redirected:
+  ret i32 0
+}
+
+define internal void @set_return_addr(i64* %addr) {
+; CHECK-LABEL: @set_return_addr(
+; CHECK-NEXT:    unreachable
+;
+  %addr.addr = alloca i64*, i32 0, align 8
+  store i64* %addr, i64** %addr.addr, align 8
+  ret void
+}