rocm: update rocm example to check for errors

This commit also adds a check for the number of devices available on the system. Signed-off-by: Daniel Bevenius <[email protected]>
danbev · Jun 15, 2024 · bf31113 · bf31113
1 parent cbfc2a4
commit bf31113
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/gpu/rocm/Makefile b/gpu/rocm/Makefile
@@ -4,4 +4,4 @@ helloworld: src/helloworld.cpp
 
 .PHONY: clean
 clean:
-	${RM} hello-world
+	${RM} helloworld
diff --git a/gpu/rocm/src/helloworld.cpp b/gpu/rocm/src/helloworld.cpp
@@ -6,21 +6,33 @@ __global__ void gpuKernel(int *flag) {
     *flag = 1; // Set flag to indicate the GPU kernel was executed
 }
 
+#define hipCheckError(ans) { hipAssert((ans), __FILE__, __LINE__); }
+inline void hipAssert(hipError_t code, const char *file, int line, bool abort=true) {
+   if (code != hipSuccess) {
+      fprintf(stderr,"HIP Error: %s %s %d\n", hipGetErrorString(code), file, line);
+      if (abort) exit(code);
+   }
+}
+
 int main() {
+    int numDevices;
+    hipGetDeviceCount(&numDevices);
+    std::cout << "Number of devices: " << numDevices << std::endl;
     int *d_flag; // Device pointer
     int h_flag = 0; // Host flag
 
     // Allocate memory on the GPU
-    hipMalloc((void **)&d_flag, sizeof(int));
+    hipCheckError(hipMalloc((void **)&d_flag, sizeof(int)));
 
     // Initialize flag on GPU
-    hipMemcpy(d_flag, &h_flag, sizeof(int), hipMemcpyHostToDevice);
+    hipCheckError(hipMemcpy(d_flag, &h_flag, sizeof(int), hipMemcpyHostToDevice));
 
     // Launch kernel (1 block, 1 thread)
     hipLaunchKernelGGL(gpuKernel, dim3(1), dim3(1), 0, 0, d_flag);
+    hipDeviceSynchronize();
 
     // Copy back the flag to host
-    hipMemcpy(&h_flag, d_flag, sizeof(int), hipMemcpyDeviceToHost);
+    hipCheckError(hipMemcpy(&h_flag, d_flag, sizeof(int), hipMemcpyDeviceToHost));
 
     // Check if the GPU kernel was executed
     if (h_flag == 1) {