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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/213635.diff


3 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp (+9-2) 
- (modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+9) 
- (added) llvm/test/CodeGen/AMDGPU/llvm.amdgcn.ballot.i32.wave64.err.ll (+15) 


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

``````````

</details>


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