[llvm] [AMDGPU] Fix llvm.amdgcn.ballot with return width != wavefront size (PR #211493)
Arseniy Obolenskiy via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 23 01:23:51 PDT 2026
https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/211493
Before wave mask was emitted directly in the requested return type, which failed to select for i32 ballots on wave64 (and vice versa)
Compute the mask at the wavefront width and then zext or trunc it to the result type
>From ffb461f4dbf3c2306fbcaa8e4da97bba64c2a6a0 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 23 Jul 2026 10:21:25 +0200
Subject: [PATCH] [AMDGPU] Fix llvm.amdgcn.ballot with return width !=
wavefront size
Before wave mask was emitted directly in the requested return type, which failed to select for i32 ballots on wave64 (and vice versa)
---
.../AMDGPU/AMDGPUInstructionSelector.cpp | 21 ++++--
llvm/lib/Target/AMDGPU/SIISelLowering.cpp | 16 ++++-
.../AMDGPU/llvm.amdgcn.ballot.i32.wave64.ll | 55 ++++++++++++++++
.../AMDGPU/llvm.amdgcn.ballot.i64.wave32.ll | 65 ++++++++++++-------
.../CodeGen/AMDGPU/nor-divergent-lanemask.ll | 7 +-
5 files changed, 127 insertions(+), 37 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.ll
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
index 22b8b10554928..585672764b450 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
@@ -1729,16 +1729,19 @@ bool AMDGPUInstructionSelector::selectBallot(MachineInstr &I) const {
const unsigned BallotSize = MRI->getType(DstReg).getSizeInBits();
const unsigned WaveSize = STI.getWavefrontSize();
- // In the common case, the return type matches the wave size.
- // However we also support emitting i64 ballots in wave32 mode.
- if (BallotSize != WaveSize && (BallotSize != 64 || WaveSize != 32))
+ // The mask is always computed at the wavefront width and then truncated or
+ // zero-extended to the requested return type. Only i32 and i64 ballots are
+ // supported.
+ if ((BallotSize != 32 && BallotSize != 64) ||
+ (WaveSize != 32 && WaveSize != 64))
return false;
std::optional<ValueAndVReg> Arg =
getIConstantVRegValWithLookThrough(SrcReg, *MRI);
Register Dst = DstReg;
- // i64 ballot on Wave32: new Dst(i32) for WaveSize ballot.
+ // When the return type does not match the wave size, compute the mask into a
+ // temporary of the wavefront width first.
if (BallotSize != WaveSize) {
Dst = MRI->createVirtualRegister(TRI.getBoolRC());
}
@@ -1773,8 +1776,8 @@ bool AMDGPUInstructionSelector::selectBallot(MachineInstr &I) const {
}
}
- // i64 ballot on Wave32: zero-extend i32 ballot to i64.
- if (BallotSize != WaveSize) {
+ if (BallotSize > WaveSize) {
+ // i64 ballot on Wave32: zero-extend the i32 mask to i64.
Register HiReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
BuildMI(*BB, &I, DL, TII.get(AMDGPU::S_MOV_B32), HiReg).addImm(0);
BuildMI(*BB, &I, DL, TII.get(AMDGPU::REG_SEQUENCE), DstReg)
@@ -1782,6 +1785,12 @@ bool AMDGPUInstructionSelector::selectBallot(MachineInstr &I) const {
.addImm(AMDGPU::sub0)
.addReg(HiReg)
.addImm(AMDGPU::sub1);
+ } else if (BallotSize < WaveSize) {
+ // i32 ballot on Wave64: truncate the i64 mask to its low i32.
+ if (!RBI.constrainGenericRegister(DstReg, AMDGPU::SReg_32RegClass, *MRI))
+ return false;
+ BuildMI(*BB, &I, DL, TII.get(AMDGPU::COPY), DstReg)
+ .addReg(Dst, {}, AMDGPU::sub0);
}
I.eraseFromParent();
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index f37e531d39648..55479398fe083 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -7940,6 +7940,9 @@ static SDValue lowerBALLOTIntrinsic(const SITargetLowering &TLI, SDNode *N,
SDValue Src = N->getOperand(1);
SDLoc SL(N);
+ unsigned WavefrontSize = TLI.getSubtarget()->getWavefrontSize();
+ EVT CCVT = EVT::getIntegerVT(*DAG.getContext(), WavefrontSize);
+
if (Src.getOpcode() == ISD::SETCC) {
SDValue Op0 = Src.getOperand(0);
SDValue Op1 = Src.getOperand(1);
@@ -7949,7 +7952,11 @@ static SDValue lowerBALLOTIntrinsic(const SITargetLowering &TLI, SDNode *N,
Op1 = DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, Op1);
}
// (ballot (ISD::SETCC ...)) -> (AMDGPUISD::SETCC ...)
- return DAG.getNode(AMDGPUISD::SETCC, SL, VT, Op0, Op1, Src.getOperand(2));
+ SDValue SetCC = DAG.getNode(AMDGPUISD::SETCC, SL, CCVT, Op0, Op1,
+ Src.getOperand(2));
+ if (VT.bitsEq(CCVT))
+ return SetCC;
+ return DAG.getZExtOrTrunc(SetCC, SL, VT);
}
if (const ConstantSDNode *Arg = dyn_cast<ConstantSDNode>(Src)) {
// (ballot 0) -> 0
@@ -7972,9 +7979,12 @@ static SDValue lowerBALLOTIntrinsic(const SITargetLowering &TLI, SDNode *N,
// (ballot (i1 $src)) -> (AMDGPUISD::SETCC (i32 (zext $src)) (i32 0)
// ISD::SETNE)
- return DAG.getNode(
- AMDGPUISD::SETCC, SL, VT, DAG.getZExtOrTrunc(Src, SL, MVT::i32),
+ SDValue SetCC = DAG.getNode(
+ AMDGPUISD::SETCC, SL, CCVT, DAG.getZExtOrTrunc(Src, SL, MVT::i32),
DAG.getConstant(0, SL, MVT::i32), DAG.getCondCode(ISD::SETNE));
+ if (VT.bitsEq(CCVT))
+ return SetCC;
+ return DAG.getZExtOrTrunc(SetCC, SL, VT);
}
static SDValue emitRemovedIntrinsicError(SelectionDAG &DAG, const SDLoc &DL,
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.ll
new file mode 100644
index 0000000000000..3ba1f58ee0f78
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgpu9.00 -global-isel=0 < %s | FileCheck %s --check-prefixes=CHECK,DAGISEL
+; RUN: llc -mtriple=amdgpu9.00 -global-isel=1 < %s | FileCheck %s --check-prefixes=CHECK,GISEL
+
+declare i32 @llvm.amdgcn.ballot.i32(i1)
+
+; Test ballot(0)
+
+define amdgpu_cs i32 @constant_false() {
+; CHECK-LABEL: constant_false:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_mov_b32 s0, 0
+; CHECK-NEXT: ; return to shader part epilog
+ %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 0)
+ ret i32 %ballot
+}
+
+; Test ballot(1)
+
+define amdgpu_cs i32 @constant_true() {
+; CHECK-LABEL: constant_true:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_mov_b32 s0, exec_lo
+; CHECK-NEXT: ; return to shader part epilog
+ %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 1)
+ ret i32 %ballot
+}
+
+; Test ballot of a non-comparison operation
+
+define amdgpu_cs i32 @non_compare(i32 %x) {
+; CHECK-LABEL: non_compare:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: v_and_b32_e32 v0, 1, v0
+; CHECK-NEXT: v_cmp_ne_u32_e64 s[0:1], 0, v0
+; CHECK-NEXT: ; return to shader part epilog
+ %trunc = trunc i32 %x to i1
+ %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 %trunc)
+ ret i32 %ballot
+}
+
+; Test ballot of comparisons
+
+define amdgpu_cs i32 @compare_ints(i32 %x, i32 %y) {
+; CHECK-LABEL: compare_ints:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: v_cmp_eq_u32_e64 s[0:1], v0, v1
+; CHECK-NEXT: ; return to shader part epilog
+ %cmp = icmp eq i32 %x, %y
+ %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 %cmp)
+ ret i32 %ballot
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; DAGISEL: {{.*}}
+; GISEL: {{.*}}
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i64.wave32.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i64.wave32.ll
index cc942e6b802bb..cd6928795d63a 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i64.wave32.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i64.wave32.ll
@@ -95,12 +95,19 @@ define amdgpu_cs i64 @compare_floats(float %x, float %y) {
}
define amdgpu_cs i64 @ctpop_of_ballot(float %x, float %y) {
-; CHECK-LABEL: ctpop_of_ballot:
-; CHECK: ; %bb.0:
-; CHECK-NEXT: v_cmp_gt_f32_e64 s0, v0, v1
-; CHECK-NEXT: s_mov_b32 s1, 0
-; CHECK-NEXT: s_bcnt1_i32_b64 s0, s[0:1]
-; CHECK-NEXT: ; return to shader part epilog
+; DAGISEL-LABEL: ctpop_of_ballot:
+; DAGISEL: ; %bb.0:
+; DAGISEL-NEXT: v_cmp_gt_f32_e32 vcc_lo, v0, v1
+; DAGISEL-NEXT: s_mov_b32 s1, 0
+; DAGISEL-NEXT: s_bcnt1_i32_b32 s0, vcc_lo
+; DAGISEL-NEXT: ; return to shader part epilog
+;
+; GISEL-LABEL: ctpop_of_ballot:
+; GISEL: ; %bb.0:
+; GISEL-NEXT: v_cmp_gt_f32_e64 s0, v0, v1
+; GISEL-NEXT: s_mov_b32 s1, 0
+; GISEL-NEXT: s_bcnt1_i32_b64 s0, s[0:1]
+; GISEL-NEXT: ; return to shader part epilog
%cmp = fcmp ogt float %x, %y
%ballot = call i64 @llvm.amdgcn.ballot.i64(i1 %cmp)
%bcnt = call i64 @llvm.ctpop.i64(i64 %ballot)
@@ -108,19 +115,31 @@ define amdgpu_cs i64 @ctpop_of_ballot(float %x, float %y) {
}
define amdgpu_cs i32 @branch_divergent_ballot64_ne_zero_compare(i32 %v) {
-; CHECK-LABEL: branch_divergent_ballot64_ne_zero_compare:
-; CHECK: ; %bb.0:
-; CHECK-NEXT: v_cmp_gt_u32_e64 s0, 12, v0
-; CHECK-NEXT: s_mov_b32 s1, 0
-; CHECK-NEXT: s_cmp_eq_u64 s[0:1], 0
-; CHECK-NEXT: s_cbranch_scc1 .LBB7_2
-; CHECK-NEXT: ; %bb.1: ; %true
-; CHECK-NEXT: s_mov_b32 s0, 42
-; CHECK-NEXT: s_branch .LBB7_3
-; CHECK-NEXT: .LBB7_2: ; %false
-; CHECK-NEXT: s_mov_b32 s0, 33
-; CHECK-NEXT: s_branch .LBB7_3
-; CHECK-NEXT: .LBB7_3:
+; DAGISEL-LABEL: branch_divergent_ballot64_ne_zero_compare:
+; DAGISEL: ; %bb.0:
+; DAGISEL-NEXT: v_cmp_gt_u32_e32 vcc_lo, 12, v0
+; DAGISEL-NEXT: s_cbranch_vccz .LBB7_2
+; DAGISEL-NEXT: ; %bb.1: ; %true
+; DAGISEL-NEXT: s_mov_b32 s0, 42
+; DAGISEL-NEXT: s_branch .LBB7_3
+; DAGISEL-NEXT: .LBB7_2: ; %false
+; DAGISEL-NEXT: s_mov_b32 s0, 33
+; DAGISEL-NEXT: s_branch .LBB7_3
+; DAGISEL-NEXT: .LBB7_3:
+;
+; GISEL-LABEL: branch_divergent_ballot64_ne_zero_compare:
+; GISEL: ; %bb.0:
+; GISEL-NEXT: v_cmp_gt_u32_e64 s0, 12, v0
+; GISEL-NEXT: s_mov_b32 s1, 0
+; GISEL-NEXT: s_cmp_eq_u64 s[0:1], 0
+; GISEL-NEXT: s_cbranch_scc1 .LBB7_2
+; GISEL-NEXT: ; %bb.1: ; %true
+; GISEL-NEXT: s_mov_b32 s0, 42
+; GISEL-NEXT: s_branch .LBB7_3
+; GISEL-NEXT: .LBB7_2: ; %false
+; GISEL-NEXT: s_mov_b32 s0, 33
+; GISEL-NEXT: s_branch .LBB7_3
+; GISEL-NEXT: .LBB7_3:
%c = icmp ult i32 %v, 12
%ballot = call i64 @llvm.amdgcn.ballot.i64(i1 %c)
%ballot_ne_zero = icmp ne i64 %ballot, 0
@@ -136,12 +155,8 @@ define amdgpu_cs i32 @branch_divergent_ballot64_ne_zero_and(i32 %v1, i32 %v2) {
; DAGISEL: ; %bb.0:
; DAGISEL-NEXT: v_cmp_gt_u32_e32 vcc_lo, 12, v0
; DAGISEL-NEXT: v_cmp_lt_u32_e64 s0, 34, v1
-; DAGISEL-NEXT: s_mov_b32 s1, 0
-; DAGISEL-NEXT: s_and_b32 s0, vcc_lo, s0
-; DAGISEL-NEXT: v_cndmask_b32_e64 v0, 0, 1, s0
-; DAGISEL-NEXT: v_cmp_ne_u32_e64 s0, 0, v0
-; DAGISEL-NEXT: s_cmp_eq_u64 s[0:1], 0
-; DAGISEL-NEXT: s_cbranch_scc1 .LBB8_2
+; DAGISEL-NEXT: s_and_b32 vcc_lo, vcc_lo, s0
+; DAGISEL-NEXT: s_cbranch_vccz .LBB8_2
; DAGISEL-NEXT: ; %bb.1: ; %true
; DAGISEL-NEXT: s_mov_b32 s0, 42
; DAGISEL-NEXT: s_branch .LBB8_3
diff --git a/llvm/test/CodeGen/AMDGPU/nor-divergent-lanemask.ll b/llvm/test/CodeGen/AMDGPU/nor-divergent-lanemask.ll
index 3351dffd0d3b4..dca697eb9d952 100644
--- a/llvm/test/CodeGen/AMDGPU/nor-divergent-lanemask.ll
+++ b/llvm/test/CodeGen/AMDGPU/nor-divergent-lanemask.ll
@@ -78,15 +78,16 @@ define amdgpu_ps i64 @test_or_two_uses(i64 inreg %a, i64 inreg %b) {
; SDAG-W32-LABEL: test_or_two_uses:
; SDAG-W32: ; %bb.0:
; SDAG-W32-NEXT: s_or_b32 s0, s0, s2
-; SDAG-W32-NEXT: s_mov_b32 s3, 0
+; SDAG-W32-NEXT: s_mov_b32 s1, 0
; SDAG-W32-NEXT: v_cndmask_b32_e64 v0, 0, 1, s0
; SDAG-W32-NEXT: s_xor_b32 s0, s0, -1
-; SDAG-W32-NEXT: s_mov_b32 s1, s3
+; SDAG-W32-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(VALU_DEP_2)
; SDAG-W32-NEXT: v_cndmask_b32_e64 v1, 0, 1, s0
-; SDAG-W32-NEXT: s_delay_alu instid0(VALU_DEP_2) | instskip(NEXT) | instid1(VALU_DEP_2)
; SDAG-W32-NEXT: v_cmp_ne_u32_e64 s0, 0, v0
+; SDAG-W32-NEXT: s_delay_alu instid0(VALU_DEP_2)
; SDAG-W32-NEXT: v_cmp_ne_u32_e64 s2, 0, v1
; SDAG-W32-NEXT: s_and_b64 s[0:1], s[2:3], s[0:1]
+; SDAG-W32-NEXT: s_mov_b32 s1, 0
; SDAG-W32-NEXT: ; return to shader part epilog
;
; GISEL-W32-LABEL: test_or_two_uses:
More information about the llvm-commits
mailing list