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

[via llvm-commits]

Author: Arseniy Obolenskiy
Date: 2026-08-03T17:52:50+02:00
New Revision: fe41b737af5229a22313ac18b760fa5984c4afb1

URL: https://github.com/llvm/llvm-project/commit/fe41b737af5229a22313ac18b760fa5984c4afb1
DIFF: https://github.com/llvm/llvm-project/commit/fe41b737af5229a22313ac18b760fa5984c4afb1.diff

LOG: Reland "[AMDGPU] Fix llvm.amdgcn.ballot with return width != wavefront size" (#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.

Added: 
    llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp
    llvm/lib/Target/AMDGPU/SIISelLowering.cpp

Removed: 
    


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