[clang] [clang][opencl] Allow passing all zeros to reqd_work_group_size (PR #131543)

[Alexander Shaposhnikov via cfe-commits]

https://github.com/alexander-shaposhnikov updated https://github.com/llvm/llvm-project/pull/131543

>From e1781606a72e86a13b032928f35403a2859f38b2 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <ashaposhnikov at google.com>
Date: Sun, 16 Mar 2025 22:27:05 +0000
Subject: [PATCH] [clang][opencl] Allow passing all zeros to
 reqd_work_group_size

---
 clang/lib/Sema/SemaDeclAttr.cpp               | 14 ++++++++++----
 clang/test/CodeGenCUDASPIRV/spirv-attrs.cu    |  6 ++++++
 clang/test/SemaOpenCL/invalid-kernel-attrs.cl |  2 ++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index bc858c63f69b6..769f5316ed934 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -2923,10 +2923,16 @@ static void handleWorkGroupSize(Sema &S, Decl *D, const ParsedAttr &AL) {
     if (!S.checkUInt32Argument(AL, E, WGSize[i], i,
                                /*StrictlyUnsigned=*/true))
       return;
-    if (WGSize[i] == 0) {
-      S.Diag(AL.getLoc(), diag::err_attribute_argument_is_zero)
-          << AL << E->getSourceRange();
-      return;
+  }
+
+  if (!llvm::all_of(WGSize, [](uint32_t Size) { return Size == 0; })) {
+    for (unsigned i = 0; i < 3; ++i) {
+      const Expr *E = AL.getArgAsExpr(i);
+      if (WGSize[i] == 0) {
+        S.Diag(AL.getLoc(), diag::err_attribute_argument_is_zero)
+            << AL << E->getSourceRange();
+        return;
+      }
     }
   }
 
diff --git a/clang/test/CodeGenCUDASPIRV/spirv-attrs.cu b/clang/test/CodeGenCUDASPIRV/spirv-attrs.cu
index 466aee00717a0..727e0e233329c 100644
--- a/clang/test/CodeGenCUDASPIRV/spirv-attrs.cu
+++ b/clang/test/CodeGenCUDASPIRV/spirv-attrs.cu
@@ -3,6 +3,9 @@
 
 #define __global__ __attribute__((global))
 
+__attribute__((reqd_work_group_size(0, 0, 0)))
+__global__ void reqd_work_group_size_0_0_0() {}
+
 __attribute__((reqd_work_group_size(128, 1, 1)))
 __global__ void reqd_work_group_size_128_1_1() {}
 
@@ -15,6 +18,8 @@ __global__ void vec_type_hint_int() {}
 __attribute__((intel_reqd_sub_group_size(64)))
 __global__ void intel_reqd_sub_group_size_64() {}
 
+
+// CHECK: define spir_kernel void @_Z26reqd_work_group_size_0_0_0v() #[[ATTR:[0-9]+]] !reqd_work_group_size ![[WG_SIZE_ZEROS:[0-9]+]]
 // CHECK: define spir_kernel void @_Z28reqd_work_group_size_128_1_1v() #[[ATTR:[0-9]+]] !reqd_work_group_size ![[WG_SIZE:[0-9]+]]
 // CHECK: define spir_kernel void @_Z26work_group_size_hint_2_2_2v() #[[ATTR]] !work_group_size_hint ![[WG_HINT:[0-9]+]]
 // CHECK: define spir_kernel void @_Z17vec_type_hint_intv() #[[ATTR]] !vec_type_hint ![[VEC_HINT:[0-9]+]]
@@ -22,6 +27,7 @@ __global__ void intel_reqd_sub_group_size_64() {}
 
 // CHECK: attributes #[[ATTR]] = { {{.*}} }
 
+// CHECK: ![[WG_SIZE_ZEROS]] = !{i32 0, i32 0, i32 0}
 // CHECK: ![[WG_SIZE]] = !{i32 128, i32 1, i32 1}
 // CHECK: ![[WG_HINT]] = !{i32 2, i32 2, i32 2}
 // CHECK: ![[VEC_HINT]] = !{i32 undef, i32 1}
diff --git a/clang/test/SemaOpenCL/invalid-kernel-attrs.cl b/clang/test/SemaOpenCL/invalid-kernel-attrs.cl
index 0883379601ef2..e913e363ef4a1 100644
--- a/clang/test/SemaOpenCL/invalid-kernel-attrs.cl
+++ b/clang/test/SemaOpenCL/invalid-kernel-attrs.cl
@@ -44,3 +44,5 @@ __kernel __attribute__((reqd_work_group_size(8,16,-32))) void neg2(){} // expect
 // 4294967294 is a negative integer if treated as signed.
 // Should compile successfully, since we expect an unsigned.
 __kernel __attribute__((reqd_work_group_size(8,16,4294967294))) void ok1(){}
+
+__kernel __attribute__((reqd_work_group_size(0,0,0))) void ok_zeros(){}