[JumpThreading] Do not unfold select if block has address taken and used #135106
Conversation
|
@llvm/pr-subscribers-llvm-transforms Author: Weihang Fan (weihangf-apple) ChangesIn 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:
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}
|
|
Hello @hstk30-hw , it looks like @xortator is inactive, would it be possible for you to suggest another reviewer? |
| 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) |
There was a problem hiding this comment.
Do you have a reproducer that is not UB? This is jumping to a block that is not listed in the indirectbr target list.
There was a problem hiding this comment.
Thank you for your feedback. I have changed the test in a way that I believe does not have UB anymore. For reference, this is the IR produced by JumpThreading before the fix:
define i64 @ba_unfold(i1 %cond, i32 %arg) {
B1:
br i1 %cond, label %B2, label %B5
B2: ; preds = %B1
br label %while.body
B3: ; preds = %while.body
%new_i = add nuw nsw i32 1, %i
%exitcond = icmp eq i32 %new_i, 5
br i1 %exitcond, label %B6, label %while.body
while.body: ; preds = %B3, %B2
%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 %B5
B5: ; preds = %while.body, %B1
ret i64 0
B6: ; preds = %B3
ret i64 ptrtoint (ptr inttoptr (i32 1 to ptr) to i64)
}
In B6, returning 1 is not correct as it is not the block address of B4.
There was a problem hiding this comment.
This test doesn't have UB anymore, but I believe the transform is correct:
This value only has defined behavior when used as an operand to the ‘indirectbr’ or for comparisons against null. Pointer equality tests between labels addresses results in undefined behavior — though, again, comparison against null is ok, and no label is equal to the null pointer. This may be passed around as an opaque pointer sized value as long as the bits are not inspected. This allows ptrtoint and arithmetic to be performed on these values so long as the original value is reconstituted before the indirectbr instruction.
Note that the use of blockaddress does not prevent the removal of the block being referenced, if it's not an indirectbr target.
Can you share any context on how you end up with this blockaddress in the first place? This suggests you are using block labels for something other than computed goto, which is not supported. Is this a ({ __label__ __here; __here: (unsigned long)&&__here; }) "get the instruction pointer" style usage perhaps, or is it something else?
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.