[CodeGen] Check BlockAddress users before marking block as taken

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->hasZeroLiveUses())
+          MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+    }
 
     if (!HasMustTailInVarArgFn)
       HasMustTailInVarArgFn = checkForMustTailInVarArgFn(IsVarArg, BB);

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->hasZeroLiveUses())
+          MBB->setAddressTakenIRBlock(const_cast<BasicBlock *>(&BB));
+    }
 
     // Mark landing pad blocks.
     if (BB.isEHPad())

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
+}

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.
+
+@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
+}

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