[llvm] Reland "[AMDGPU] Fix llvm.amdgcn.ballot with return width != wavefront size" (PR #213635)

[Arseniy Obolenskiy via llvm-commits]

https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/213635

Reverts https://github.com/llvm/llvm-project/pull/212628

This relands #211493, which was reverted because ockl_dm_alloc/ockl_dm_dealloc in device-libs emit an i32 ballot on wave64, which GlobalISel cannot select (one bit per lane doesn't fit). [#212813](https://github.com/llvm/llvm-project/pull/212813) widens the clang ballot builtins to the wavefront size so a narrower-than-wave ballot is no longer emitted, fixing the root cause.

>From 1d1b641bd9d300ea667bdd092d69866d5ec7b82b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Mon, 3 Aug 2026 11:50:44 +0200
Subject: [PATCH] Reland "[AMDGPU] Fix llvm.amdgcn.ballot with return width !=
 wavefront size"

Reverts https://github.com/llvm/llvm-project/pull/212628

This relands #211493, which was reverted because ockl_dm_alloc/ockl_dm_dealloc
in device-libs emit an i32 ballot on wave64, which GlobalISel cannot select
(one bit per lane doesn't fit). #212813 widens the clang ballot builtins to
the wavefront size so a narrower-than-wave ballot is no longer emitted,
fixing the root cause.
---
 .../Target/AMDGPU/AMDGPUInstructionSelector.cpp   | 11 +++++++++--
 llvm/lib/Target/AMDGPU/SIISelLowering.cpp         |  9 +++++++++
 .../AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll   | 15 +++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
index 22b8b10554928..0ddb7f771b6c9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
@@ -1729,8 +1729,15 @@ 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) {
+    const Function &Fn = MF->getFunction();
+    Fn.getContext().diagnose(DiagnosticInfoUnsupported(
+        Fn, "ballot return type is narrower than the wavefront size", DL));
+    BuildMI(*BB, &I, DL, TII.get(AMDGPU::IMPLICIT_DEF), DstReg);
+    I.eraseFromParent();
+    return true;
+  }
+
   if (BallotSize != WaveSize && (BallotSize != 64 || WaveSize != 32))
     return false;
 
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 612342b159583..40c662a89a4c4 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -7957,6 +7957,15 @@ static SDValue lowerBALLOTIntrinsic(const SITargetLowering &TLI, SDNode *N,
   SDValue Src = N->getOperand(1);
   SDLoc SL(N);
 
+  unsigned WavefrontSize = TLI.getSubtarget()->getWavefrontSize();
+  if (VT.getScalarSizeInBits() < WavefrontSize) {
+    DAG.getContext()->diagnose(DiagnosticInfoUnsupported(
+        DAG.getMachineFunction().getFunction(),
+        "ballot return type is narrower than the wavefront size",
+        SL.getDebugLoc()));
+    return DAG.getPOISON(VT);
+  }
+
   if (Src.getOpcode() == ISD::SETCC) {
     SDValue Op0 = Src.getOperand(0);
     SDValue Op1 = Src.getOperand(1);
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll
new file mode 100644
index 0000000000000..54400baf4affe
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll
@@ -0,0 +1,15 @@
+; RUN: not llc -global-isel=0 -mtriple=amdgpu9.00 -filetype=null %s 2>&1 | FileCheck %s
+; RUN: not llc -global-isel=1 -mtriple=amdgpu9.00 -filetype=null %s 2>&1 | FileCheck %s
+
+; An i32 ballot on a wave64 target cannot represent one bit per lane, so
+; both SelectionDAG and GlobalISel must refuse to lower it instead of
+; dropping the mask bits of the high lanes.
+
+declare i32 @llvm.amdgcn.ballot.i32(i1)
+
+; CHECK: error: {{.*}}ballot return type is narrower than the wavefront size
+define amdgpu_cs i32 @ballot_i32_wave64(i32 %x, i32 %y) {
+  %cmp = icmp eq i32 %x, %y
+  %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 %cmp)
+  ret i32 %ballot
+}