[llvm] [JumpThreading] Do not thread if block has address taken and used (PR #135106)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Weihang Fan (weihangf-apple)

<details>
<summary>Changes</summary>

In an internal test case, JumpThreadingPass threaded across a basic block that has its address taken, resulting in the block being removed and accidentally invalidating a blockaddress. This caused SimplifyCFG to insert an unreachable instruction which ultimately caused a crash at runtime.

The issue in the case was that the select instruction was unfolded, causing two basic blocks to be created, which meant that the blockaddress of the original block is now ambiguous.

This PR guards against unfolding selects in basic blocks that have their address taken and used and adds a unit test.

---
Full diff: https://github.com/llvm/llvm-project/pull/135106.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/JumpThreading.cpp (+5) 
- (modified) llvm/test/Transforms/JumpThreading/select.ll (+46) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index a518e02d762f6..0c189ae769e88 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -2933,6 +2933,11 @@ bool JumpThreadingPass::tryToUnfoldSelectInCurrBB(BasicBlock *BB) {
   if (LoopHeaders.count(BB))
     return false;
 
+  // If a block has its address taken, it would break the semantics of its block
+  // address. To be safe, don't thread the edge.
+  if (hasAddressTakenAndUsed(BB))
+    return false;
+
   for (BasicBlock::iterator BI = BB->begin();
        PHINode *PN = dyn_cast<PHINode>(BI); ++BI) {
     // Look for a Phi having at least one constant incoming value.
diff --git a/llvm/test/Transforms/JumpThreading/select.ll b/llvm/test/Transforms/JumpThreading/select.ll
index 4ec55a66bb8ac..1774ce9594337 100644
--- a/llvm/test/Transforms/JumpThreading/select.ll
+++ b/llvm/test/Transforms/JumpThreading/select.ll
@@ -665,6 +665,52 @@ if.end:
   ret i32 %v1
 }
 
+; Do not unfold select when parent block address is taken and used
+define i32 @ba_unfold(i1 %cond, i32 %arg) {
+; CHECK-LABEL: @ba_unfold(
+; CHECK: B4:
+; CHECK-NEXT:   [[P:%.*]] = phi i1 [ true, [[B1:%.*]] ], [ [[COND_NOT:%.*]], [[LOOPEXIT:%.*]] ]
+; CHECK-NEXT:   [[OFFSET:%.*]] = select i1 [[P]], i64 8, i64 16
+; CHECK-NEXT:   [[ADDR_I:%.*]] = add i64 [[OFFSET]], ptrtoint (ptr blockaddress(@ba_unfold, %B4) to i64)
+;
+entry:
+  br label %B1
+
+B1:
+  br i1 %cond, label %B2, label %B4
+
+B2:
+  br label %while.body
+
+B3:
+  %new_i = add nuw nsw i32 1, %i
+  %exitcond = icmp eq i32 %new_i, 5
+  br i1 %exitcond, label %loopexit, label %while.body
+
+while.body:
+  %i = phi i32 [ 0, %B2 ], [ %new_i, %B3 ]
+  %xor = xor i32 %i, 16
+  %cmp = icmp eq i32 %xor, %arg
+  br i1 %cmp, label %B3, label %loopexit
+
+loopexit:
+  %cond.not = xor i1 %cmp, true
+  br label %B4
+
+B4:
+  %p = phi i1 [true, %B1], [ %cond.not, %loopexit ]
+  %offset = select i1 %p, i64 8, i64 16
+  %addr_i = add i64 %offset, ptrtoint (ptr blockaddress(@ba_unfold, %B4) to i64)
+  %addr = inttoptr i64 %addr_i to ptr
+  indirectbr ptr %addr, [label %B5, label %B6]
+
+B5:
+  br label %B6
+
+B6:
+  ret i32 0
+}
+
 ; branch_weights overflowing uint32_t
 !0 = !{!"branch_weights", i64 1073741824, i64 3221225472}
 !1 = !{!"function_entry_count", i64 1984}

``````````

</details>


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