[llvm] [Offload] Set size correctly in olLaunchKernel cts test (PR #142398)

[Ross Brunton via llvm-commits]

https://github.com/RossBrunton created https://github.com/llvm/llvm-project/pull/142398

It was previously not scaled by `sizeof(uint32_t)`.


>From d5121ce17c23a3b1e91d7995273e1c892d54e2c9 Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Mon, 2 Jun 2025 15:01:58 +0100
Subject: [PATCH] [Offload] Set size correctly in olLaunchKernel cts test

It was previously not scaled by `sizeof(uint32_t)`.
---
 offload/unittests/OffloadAPI/device_code/foo.c         |  3 ++-
 offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp | 10 ++++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/offload/unittests/OffloadAPI/device_code/foo.c b/offload/unittests/OffloadAPI/device_code/foo.c
index 5bc893961d49f..83cdc53cddd8d 100644
--- a/offload/unittests/OffloadAPI/device_code/foo.c
+++ b/offload/unittests/OffloadAPI/device_code/foo.c
@@ -1,5 +1,6 @@
 #include <gpuintrin.h>
+#include <stdint.h>
 
-__gpu_kernel void foo(int *out) {
+__gpu_kernel void foo(uint32_t *out) {
   out[__gpu_thread_id(0)] = __gpu_thread_id(0);
 }
diff --git a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
index 20462e22fd73f..e17d094cfa612 100644
--- a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
+++ b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
@@ -47,7 +47,8 @@ OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olLaunchKernelTest);
 
 TEST_P(olLaunchKernelTest, Success) {
   void *Mem;
-  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, 64, &Mem));
+  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED,
+                            LaunchArgs.GroupSizeX * sizeof(uint32_t), &Mem));
   struct {
     void *Mem;
   } Args{Mem};
@@ -57,7 +58,7 @@ TEST_P(olLaunchKernelTest, Success) {
 
   ASSERT_SUCCESS(olWaitQueue(Queue));
 
-  int *Data = (int *)Mem;
+  uint32_t *Data = (uint32_t *)Mem;
   for (int i = 0; i < 64; i++) {
     ASSERT_EQ(Data[i], i);
   }
@@ -67,7 +68,8 @@ TEST_P(olLaunchKernelTest, Success) {
 
 TEST_P(olLaunchKernelTest, SuccessSynchronous) {
   void *Mem;
-  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, 64, &Mem));
+  ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED,
+                            LaunchArgs.GroupSizeX * sizeof(uint32_t), &Mem));
 
   struct {
     void *Mem;
@@ -76,7 +78,7 @@ TEST_P(olLaunchKernelTest, SuccessSynchronous) {
   ASSERT_SUCCESS(olLaunchKernel(nullptr, Device, Kernel, &Args, sizeof(Args),
                                 &LaunchArgs, nullptr));
 
-  int *Data = (int *)Mem;
+  uint32_t *Data = (uint32_t *)Mem;
   for (int i = 0; i < 64; i++) {
     ASSERT_EQ(Data[i], i);
   }