From 8d02e942a6b3d773b4c382d44d18f9f266b8a465 Mon Sep 17 00:00:00 2001
From: Michael Xie <pengda.xie@amd.com>
Date: Mon, 14 Oct 2024 15:24:37 -0700
Subject: [PATCH] SWDEV-488231 - Add memory cleanup after test

Change-Id: I0a4f6e1649c1c30e6f76243f83681a6755b5368b
---
 catch/multiproc/deviceAllocationMproc.cc | 28 ++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/catch/multiproc/deviceAllocationMproc.cc b/catch/multiproc/deviceAllocationMproc.cc
index eccbb76094..e91c0443ce 100644
--- a/catch/multiproc/deviceAllocationMproc.cc
+++ b/catch/multiproc/deviceAllocationMproc.cc
@@ -140,6 +140,8 @@ static bool testDeviceAllocMulProc(bool testmalloc) {
   childpid = fork();
   if (childpid > 0) {  // Parent
     close(fd[1]);
+    int *result_d{nullptr};
+    HIP_CHECK(hipMalloc(&result_d, sizeof(int)));
     // Allocate in parent
     if (testmalloc) {
       kerTestDeviceMalloc<<<1, 1>>>(SIZE);
@@ -150,6 +152,16 @@ static bool testDeviceAllocMulProc(bool testmalloc) {
     // Check allocated memory size
     HIP_CHECK(hipMemGetInfo(&avail, &tot));
     if ((tot - avail) < SIZE) {
+      // Clean up memory before return
+      if (testmalloc) {
+        kerTestDeviceFree<<<1, 1>>>(result_d);
+      } else {
+        kerTestDeviceDelete<<<1, 1>>>(result_d);
+      }
+      HIP_CHECK(hipDeviceSynchronize());
+      HIP_CHECK(hipFree(result_d));
+      close(fd[0]);
+      wait(NULL);
       return false;
     }
     // parent will wait to read the device cnt
@@ -164,8 +176,17 @@ static bool testDeviceAllocMulProc(bool testmalloc) {
     if ((tot - avail) < SIZE) {
       testResult = false;
     }
+    if (testmalloc) {
+      kerTestDeviceFree<<<1, 1>>>(result_d);
+    } else {
+      kerTestDeviceDelete<<<1, 1>>>(result_d);
+    }
+    HIP_CHECK(hipDeviceSynchronize());
+    HIP_CHECK(hipFree(result_d));
   } else if (!childpid) {  // Child
     // Wait for hipDeviceSetLimit() completion in parent.
+    int *result_d{nullptr};
+    HIP_CHECK(hipMalloc(&result_d, sizeof(int)));
     close(fd[0]);
     // Allocate in child
     if (testmalloc) {
@@ -185,6 +206,13 @@ static bool testDeviceAllocMulProc(bool testmalloc) {
     write(fd[1], &testResult, sizeof(testResult));
     // close the write descriptor:
     close(fd[1]);
+    if (testmalloc) {
+      kerTestDeviceFree<<<1, 1>>>(result_d);
+    } else {
+      kerTestDeviceDelete<<<1, 1>>>(result_d);
+    }
+    HIP_CHECK(hipDeviceSynchronize());
+    HIP_CHECK(hipFree(result_d));
     exit(0);
   }
   return testResult;