[clang] [clang][AMDGPU] Widen ballot for read_exec_lo/hi to wavefront size (PR #212813)

[via cfe-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-clang

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

GlobalISel cannot select a ballot narrower than the wavefront width, since it can't represent one bit per lane

Widen the ballot to the wave size and narrow the result afterwards

This is a prerequisite for relanding https://github.com/llvm/llvm-project/pull/211493 (reverted in https://github.com/llvm/llvm-project/pull/212628 to unblock buildbot) to prevent device libs side failures

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


4 Files Affected:

- (modified) clang/lib/Basic/Targets/SPIR.h (+4) 
- (modified) clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp (+11-2) 
- (modified) clang/test/CodeGenOpenCL/builtins-amdgcn-wave64.cl (+2-3) 
- (modified) clang/test/CodeGenOpenCL/builtins-amdgcn.cl (+2-3) 


``````````diff
diff --git a/clang/lib/Basic/Targets/SPIR.h b/clang/lib/Basic/Targets/SPIR.h
index a16c8ec79d2f2..7cfe927ac6463 100644
--- a/clang/lib/Basic/Targets/SPIR.h
+++ b/clang/lib/Basic/Targets/SPIR.h
@@ -453,6 +453,10 @@ class LLVM_LIBRARY_VISIBILITY SPIRV64AMDGCNTargetInfo final
   void getTargetDefines(const LangOptions &Opts,
                         MacroBuilder &Builder) const override;
 
+  const llvm::omp::GV &getGridValue() const override {
+    return llvm::omp::SPIRVGridValues;
+  }
+
   void setAuxTarget(const TargetInfo *Aux) override;
 
   void adjust(DiagnosticsEngine &Diags, LangOptions &Opts,
diff --git a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
index 72d6f536165c5..0dfb5e91f72f9 100644
--- a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
+++ b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp
@@ -275,13 +275,22 @@ Value *EmitAMDGPUGridSize(CodeGenFunction &CGF, unsigned Index) {
 
 // Generates the IR for __builtin_read_exec_*.
 // Lowers the builtin to amdgcn_ballot intrinsic.
+//
+// The ballot must be taken at the wavefront width: a ballot narrower than the
+// wave size cannot represent one bit per lane and fails to select. Request the
+// mask at the wave width and narrow it afterwards for the _lo and _hi halves.
 static Value *EmitAMDGCNBallotForExec(CodeGenFunction &CGF, const CallExpr *E,
                                       llvm::Type *RegisterType,
                                       llvm::Type *ValueType, bool isExecHi) {
   CodeGen::CGBuilderTy &Builder = CGF.Builder;
   CodeGen::CodeGenModule &CGM = CGF.CGM;
 
-  Function *F = CGM.getIntrinsic(Intrinsic::amdgcn_ballot, {RegisterType});
+  unsigned WaveSize = CGF.getTarget().getGridValue().GV_Warp_Size;
+  llvm::Type *BallotType = RegisterType;
+  if (BallotType->getIntegerBitWidth() < WaveSize)
+    BallotType = Builder.getIntNTy(WaveSize);
+
+  Function *F = CGM.getIntrinsic(Intrinsic::amdgcn_ballot, {BallotType});
   llvm::Value *Call = Builder.CreateCall(F, {Builder.getInt1(true)});
 
   if (isExecHi) {
@@ -290,7 +299,7 @@ static Value *EmitAMDGCNBallotForExec(CodeGenFunction &CGF, const CallExpr *E,
     return Rt2;
   }
 
-  return Call;
+  return Builder.CreateTrunc(Call, CGF.ConvertType(E->getType()));
 }
 
 static llvm::Value *loadTextureDescPtorAsVec8I32(CodeGenFunction &CGF,
diff --git a/clang/test/CodeGenOpenCL/builtins-amdgcn-wave64.cl b/clang/test/CodeGenOpenCL/builtins-amdgcn-wave64.cl
index ef39558f94f2c..b049031156335 100644
--- a/clang/test/CodeGenOpenCL/builtins-amdgcn-wave64.cl
+++ b/clang/test/CodeGenOpenCL/builtins-amdgcn-wave64.cl
@@ -37,13 +37,12 @@ void test_read_exec(global ulong* out) {
 }
 
 // CHECK-LABEL: @test_read_exec_lo(
-// CHECK: call i32 @llvm.amdgcn.ballot.i32(i1 true)
+// CHECK: call i64 @llvm.amdgcn.ballot.i64(i1 true)
+// CHECK: and i64 [[A:%.*]], 4294967295
 void test_read_exec_lo(global ulong* out) {
   *out = __builtin_amdgcn_read_exec_lo();
 }
 
-// CHECK: declare i32 @llvm.amdgcn.ballot.i32(i1) #[[$NOUNWIND_READONLY:[0-9]+]]
-
 // CHECK-LABEL: @test_read_exec_hi(
 // CHECK: call i64 @llvm.amdgcn.ballot.i64(i1 true)
 // CHECK: lshr i64 [[A:%.*]], 32
diff --git a/clang/test/CodeGenOpenCL/builtins-amdgcn.cl b/clang/test/CodeGenOpenCL/builtins-amdgcn.cl
index 28c420a5760f4..78a4d7af386c2 100644
--- a/clang/test/CodeGenOpenCL/builtins-amdgcn.cl
+++ b/clang/test/CodeGenOpenCL/builtins-amdgcn.cl
@@ -1048,13 +1048,12 @@ void test_read_exec(global ulong* out) {
 }
 
 // CHECK-LABEL: @test_read_exec_lo(
-// CHECK: {{.*}}call{{.*}} i32 @llvm.amdgcn.ballot.i32(i1 true)
+// CHECK: [[A:%.*]] = {{.*}}call{{.*}} i64 @llvm.amdgcn.ballot.i64(i1 true)
+// CHECK: trunc i64 [[A]] to i32
 void test_read_exec_lo(global uint* out) {
   *out = __builtin_amdgcn_read_exec_lo();
 }
 
-// CHECK: declare i32 @llvm.amdgcn.ballot.i32(i1){{.*}} #[[$NOUNWIND_READONLY_NOPOISON:[0-9]+]]
-
 // CHECK-LABEL: @test_read_exec_hi(
 // CHECK: {{.*}}call{{.*}} i64 @llvm.amdgcn.ballot.i64(i1 true)
 // CHECK: lshr i64 [[A:%.*]], 32

``````````

</details>


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