[llvm] [OpenMP][OMPX] Fix ompx_ballot_sync test (PR #94140)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/94140

The current test is not really correct because the mask is set to 0xffffffff
even if it is on an AMDGPU whose wavefront size is 64. Besides,
`__AMDGCN_WAVEFRONT_SIZE` is not set on host compilation so the verification
happens to work.


>From 4ac902da4febbbf3883d55f5784e57d3b72d332a Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Sun, 2 Jun 2024 00:58:42 -0400
Subject: [PATCH] [OpenMP][OMPX] Fix ompx_ballot_sync test

The current test is not really correct because the mask is set to 0xffffffff
even if it is on an AMDGPU whose wavefront size is 64. Besides,
`__AMDGCN_WAVEFRONT_SIZE` is not set on host compilation so the verification
happens to work.
---
 .../test/offloading/ompx_bare_ballot_sync.c   | 26 +++++++++----------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/offload/test/offloading/ompx_bare_ballot_sync.c b/offload/test/offloading/ompx_bare_ballot_sync.c
index d8e17691bf9c7..00d8ac9d91424 100644
--- a/offload/test/offloading/ompx_bare_ballot_sync.c
+++ b/offload/test/offloading/ompx_bare_ballot_sync.c
@@ -7,12 +7,6 @@
 // UNSUPPORTED: s390x-ibm-linux-gnu
 // UNSUPPORTED: s390x-ibm-linux-gnu-LTO
 
-#if defined __AMDGCN_WAVEFRONT_SIZE && __AMDGCN_WAVEFRONT_SIZE == 64
-#define MASK 0xaaaaaaaaaaaaaaaa
-#else
-#define MASK 0xaaaaaaaa
-#endif
-
 #include <assert.h>
 #include <ompx.h>
 #include <stdint.h>
@@ -23,23 +17,27 @@ int main(int argc, char *argv[]) {
   const int num_blocks = 1;
   const int block_size = 256;
   const int N = num_blocks * block_size;
-  uint64_t *data = (uint64_t *)malloc(N * sizeof(uint64_t));
-
-  for (int i = 0; i < N; ++i)
-    data[i] = i & 0x1;
+  int *res = (int *)malloc(N * sizeof(int));
 
-#pragma omp target teams ompx_bare num_teams(num_blocks) thread_limit(block_size) map(tofrom: data[0:N])
+#pragma omp target teams ompx_bare num_teams(num_blocks) thread_limit(block_size) \
+        map(from: res[0:N])
   {
     int tid = ompx_thread_id_x();
-    uint64_t mask = ompx_ballot_sync(~0U, data[tid]);
-    data[tid] += mask;
+    uint64_t mask = ompx_ballot_sync(~0LU, tid & 0x1);
+#if defined __AMDGCN_WAVEFRONT_SIZE && __AMDGCN_WAVEFRONT_SIZE == 64
+    res[tid] = mask == 0xaaaaaaaaaaaaaaaa;
+#else
+    res[tid] = mask == 0xaaaaaaaa;
+#endif
   }
 
   for (int i = 0; i < N; ++i)
-    assert(data[i] == ((i & 0x1) + MASK));
+    assert(res[i]);
 
   // CHECK: PASS
   printf("PASS\n");
 
+  free(res);
+
   return 0;
 }