[llvm] [CodeGen] Check BlockAddress users before marking block as taken (PR #174480)
Justin Stitt via llvm-commits
llvm-commits at lists.llvm.org
Tue Jan 6 09:55:21 PST 2026
https://github.com/JustinStitt updated https://github.com/llvm/llvm-project/pull/174480
>From b6263bf2b053ebef7662a430aa55301b1c21e25e Mon Sep 17 00:00:00 2001
From: Justin Stitt <justinstitt at google.com>
Date: Mon, 5 Jan 2026 11:39:46 -0800
Subject: [PATCH 1/3] [CodeGen] Check BlockAddress users before marking block
as taken
After commit 4109bac3301e ("[IR] Do not store Function inside BlockAddress"),
the `hasAddressTaken` flag on BasicBlock can become stale. When a BlockAddress
constant is created, the flag is set to true. However, if optimizations
remove all users of that BlockAddress (e.g., simplifying an `indirectbr` to
a direct branch), the BlockAddress may remain in the uniquing table with no
users, but the flag stays true.
Check that BlockAddress actually has users before marking the MBB. If
the BlockAddress exists but has no users (probably because they were
optimized away), we can skip marking the block.
The performance impact should be negligible as we still only check for
users after verifying that the block address has its address taken
somewhere, which should be false pretty often (I think?!).
Signed-off-by: Justin Stitt <justinstitt at google.com>
---
llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp | 10 ++++++++--
.../lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp | 11 +++++++++--
.../X86/speculative-load-hardening-indirect.ll | 8 ++++----
3 files changed, 21 insertions(+), 8 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 12552bce3caaa..18a6dd2c15365 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -4217,8 +4217,14 @@ bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
MBB = MF->CreateMachineBasicBlock(&BB);
MF->push_back(MBB);
- if (BB.hasAddressTaken())
- MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+ // Only mark the block if the BlockAddress actually has users. The
+ // hasAddressTaken flag may be stale if the BlockAddress was optimized away
+ // but the constant still exists in the uniquing table.
+ if (BB.hasAddressTaken()) {
+ if (BlockAddress *BA = BlockAddress::lookup(&BB))
+ if (!BA->use_empty())
+ MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+ }
if (!HasMustTailInVarArgFn)
HasMustTailInVarArgFn = checkForMustTailInVarArgFn(IsVarArg, BB);
diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
index e73743ecbc9fa..74100eeeefcdd 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -26,6 +26,7 @@
#include "llvm/CodeGen/TargetSubtargetInfo.h"
#include "llvm/CodeGen/WasmEHFuncInfo.h"
#include "llvm/CodeGen/WinEHFuncInfo.h"
+#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
@@ -277,8 +278,14 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
// Transfer the address-taken flag. This is necessary because there could
// be multiple MachineBasicBlocks corresponding to one BasicBlock, and only
// the first one should be marked.
- if (BB.hasAddressTaken())
- MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+ // Only mark the block if the BlockAddress actually has users. The
+ // hasAddressTaken flag may be stale if the BlockAddress was optimized away
+ // but the constant still exists in the uniquing table.
+ if (BB.hasAddressTaken()) {
+ if (BlockAddress *BA = BlockAddress::lookup(&BB))
+ if (!BA->use_empty())
+ MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+ }
// Mark landing pad blocks.
if (BB.isEHPad())
diff --git a/llvm/test/CodeGen/X86/speculative-load-hardening-indirect.ll b/llvm/test/CodeGen/X86/speculative-load-hardening-indirect.ll
index fd5085c8c2ac9..2092fc9cd6e00 100644
--- a/llvm/test/CodeGen/X86/speculative-load-hardening-indirect.ll
+++ b/llvm/test/CodeGen/X86/speculative-load-hardening-indirect.ll
@@ -464,34 +464,34 @@ define dso_local i32 @test_indirectbr_global(i32 %idx) nounwind {
; X64-RETPOLINE-NEXT: cmoveq %rax, %rcx
; X64-RETPOLINE-NEXT: cmpq $4, %rdx
; X64-RETPOLINE-NEXT: jne .LBB6_3
-; X64-RETPOLINE-NEXT: .Ltmp0: # Block address taken
; X64-RETPOLINE-NEXT: # %bb.6: # %bb3
; X64-RETPOLINE-NEXT: cmovneq %rax, %rcx
; X64-RETPOLINE-NEXT: shlq $47, %rcx
; X64-RETPOLINE-NEXT: movl $42, %eax
; X64-RETPOLINE-NEXT: orq %rcx, %rsp
; X64-RETPOLINE-NEXT: retq
-; X64-RETPOLINE-NEXT: .Ltmp1: # Block address taken
; X64-RETPOLINE-NEXT: .LBB6_5: # %bb2
; X64-RETPOLINE-NEXT: cmovneq %rax, %rcx
; X64-RETPOLINE-NEXT: shlq $47, %rcx
; X64-RETPOLINE-NEXT: movl $13, %eax
; X64-RETPOLINE-NEXT: orq %rcx, %rsp
; X64-RETPOLINE-NEXT: retq
-; X64-RETPOLINE-NEXT: .Ltmp2: # Block address taken
; X64-RETPOLINE-NEXT: .LBB6_4: # %bb1
; X64-RETPOLINE-NEXT: cmovneq %rax, %rcx
; X64-RETPOLINE-NEXT: shlq $47, %rcx
; X64-RETPOLINE-NEXT: movl $7, %eax
; X64-RETPOLINE-NEXT: orq %rcx, %rsp
; X64-RETPOLINE-NEXT: retq
-; X64-RETPOLINE-NEXT: .Ltmp3: # Block address taken
; X64-RETPOLINE-NEXT: .LBB6_3: # %bb0
; X64-RETPOLINE-NEXT: cmoveq %rax, %rcx
; X64-RETPOLINE-NEXT: shlq $47, %rcx
; X64-RETPOLINE-NEXT: movl $2, %eax
; X64-RETPOLINE-NEXT: orq %rcx, %rsp
; X64-RETPOLINE-NEXT: retq
+; X64-RETPOLINE-NEXT: .Ltmp0: # Address of block that was removed by CodeGen
+; X64-RETPOLINE-NEXT: .Ltmp1: # Address of block that was removed by CodeGen
+; X64-RETPOLINE-NEXT: .Ltmp2: # Address of block that was removed by CodeGen
+; X64-RETPOLINE-NEXT: .Ltmp3: # Address of block that was removed by CodeGen
entry:
%ptr = getelementptr [4 x ptr], ptr @global_blockaddrs, i32 0, i32 %idx
%a = load ptr, ptr %ptr
>From 6bb4ffbd9800534b17239c35bf14f980446196f0 Mon Sep 17 00:00:00 2001
From: Justin Stitt <justinstitt at google.com>
Date: Mon, 5 Jan 2026 15:17:36 -0800
Subject: [PATCH 2/3] add more tests
Signed-off-by: Justin Stitt <justinstitt at google.com>
---
.../blockaddress-stale-addresstaken.ll | 48 +++++++++++++++++++
.../X86/blockaddress-stale-addresstaken.ll | 39 +++++++++++++++
2 files changed, 87 insertions(+)
create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/blockaddress-stale-addresstaken.ll
create mode 100644 llvm/test/CodeGen/X86/blockaddress-stale-addresstaken.ll
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/blockaddress-stale-addresstaken.ll b/llvm/test/CodeGen/AArch64/GlobalISel/blockaddress-stale-addresstaken.ll
new file mode 100644
index 0000000000000..601645e0944f5
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/blockaddress-stale-addresstaken.ll
@@ -0,0 +1,48 @@
+; RUN: llc -O0 -mtriple=aarch64-linux-gnu -global-isel -stop-after=irtranslator %s -o - | FileCheck %s
+
+; Test that the GlobalISel IRTranslator correctly marks blocks as address-taken
+; based on whether the BlockAddress actually has users.
+
+; CHECK-LABEL: name: test_indirectbr_blockaddress
+; CHECK: G_BLOCK_ADDR blockaddress(@test_indirectbr_blockaddress, %ir-block.target)
+; CHECK: G_BLOCK_ADDR blockaddress(@test_indirectbr_blockaddress, %ir-block.other)
+; CHECK: G_BRINDIRECT
+; CHECK: bb.{{[0-9]+}}.target (ir-block-address-taken %ir-block.target):
+; CHECK: bb.{{[0-9]+}}.other (ir-block-address-taken %ir-block.other):
+define i32 @test_indirectbr_blockaddress(i32 %idx) {
+entry:
+ %targets = alloca [2 x ptr], align 8
+ %ptr0 = getelementptr [2 x ptr], ptr %targets, i64 0, i64 0
+ store ptr blockaddress(@test_indirectbr_blockaddress, %target), ptr %ptr0, align 8
+ %ptr1 = getelementptr [2 x ptr], ptr %targets, i64 0, i64 1
+ store ptr blockaddress(@test_indirectbr_blockaddress, %other), ptr %ptr1, align 8
+ %idx64 = zext i32 %idx to i64
+ %selected = getelementptr [2 x ptr], ptr %targets, i64 0, i64 %idx64
+ %dest = load ptr, ptr %selected, align 8
+ indirectbr ptr %dest, [label %target, label %other]
+
+target:
+ ret i32 42
+
+other:
+ ret i32 -1
+}
+
+; normal conditional branch (no blockaddress).
+; blocks should NOT be marked as address-taken.
+
+; CHECK-LABEL: name: test_normal_branch
+; CHECK: bb.{{[0-9]+}}.target:
+; CHECK-NOT: ir-block-address-taken
+; CHECK: bb.{{[0-9]+}}.other:
+; CHECK-NOT: ir-block-address-taken
+define i32 @test_normal_branch(i1 %cond) {
+entry:
+ br i1 %cond, label %target, label %other
+
+target:
+ ret i32 42
+
+other:
+ ret i32 -1
+}
diff --git a/llvm/test/CodeGen/X86/blockaddress-stale-addresstaken.ll b/llvm/test/CodeGen/X86/blockaddress-stale-addresstaken.ll
new file mode 100644
index 0000000000000..52adca08cb259
--- /dev/null
+++ b/llvm/test/CodeGen/X86/blockaddress-stale-addresstaken.ll
@@ -0,0 +1,39 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -mattr=+retpoline | FileCheck %s
+;
+; verify that blocks are NOT marked as "Block address taken" when the
+; BlockAddress constant has no users (was optimized away).
+;
+; With retpoline enabled, the indirectbr is replaced with direct comparisons
+; against constant integer values. The BlockAddress constants in the global
+; array become unused (constant-folded to integers), so the blocks should NOT
+; be marked as address-taken.
+
+ at targets = internal constant [4 x ptr] [
+ ptr blockaddress(@test_stale_addresstaken, %bb0),
+ ptr blockaddress(@test_stale_addresstaken, %bb1),
+ ptr blockaddress(@test_stale_addresstaken, %bb2),
+ ptr blockaddress(@test_stale_addresstaken, %bb3)
+]
+
+define i32 @test_stale_addresstaken(i32 %idx) {
+entry:
+ %ptr = getelementptr [4 x ptr], ptr @targets, i32 0, i32 %idx
+ %dest = load ptr, ptr %ptr
+ indirectbr ptr %dest, [label %bb0, label %bb1, label %bb2, label %bb3]
+
+; CHECK-LABEL: test_stale_addresstaken:
+; CHECK-NOT: Block address taken
+; CHECK-LABEL: .Lfunc_end0:
+
+bb0:
+ ret i32 0
+
+bb1:
+ ret i32 1
+
+bb2:
+ ret i32 2
+
+bb3:
+ ret i32 3
+}
>From dd814bcea4d6b083f023074b6cc812334b180a33 Mon Sep 17 00:00:00 2001
From: Justin Stitt <justinstitt at google.com>
Date: Mon, 5 Jan 2026 15:34:29 -0800
Subject: [PATCH 3/3] use hasZeroLiveUses over use_empty
Signed-off-by: Justin Stitt <justinstitt at google.com>
---
llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp | 2 +-
llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 18a6dd2c15365..e1b7e192939c9 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -4222,7 +4222,7 @@ bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
// but the constant still exists in the uniquing table.
if (BB.hasAddressTaken()) {
if (BlockAddress *BA = BlockAddress::lookup(&BB))
- if (!BA->use_empty())
+ if (!BA->hasZeroLiveUses())
MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
}
diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
index 74100eeeefcdd..dfaabae6e1f97 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -283,7 +283,7 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
// but the constant still exists in the uniquing table.
if (BB.hasAddressTaken()) {
if (BlockAddress *BA = BlockAddress::lookup(&BB))
- if (!BA->use_empty())
+ if (!BA->hasZeroLiveUses())
MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
}
More information about the llvm-commits
mailing list