From b751614cd8dee42d2c63775eb829e5f65e3abc10 Mon Sep 17 00:00:00 2001
From: "Patel, Jaydeepkumar" <Jaydeepkumar.Patel@amd.com>
Date: Tue, 27 May 2025 02:18:21 -0700
Subject: [PATCH] SWDEV-508982 - Handle hipMemPoolCreate, hipMemPoolDestory &
 hipDeviceSetMemPool during stream capture. (#398)

---
 hipamd/src/hip_mempool.cpp | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/hipamd/src/hip_mempool.cpp b/hipamd/src/hip_mempool.cpp
index ff7d164c70..114c793dbc 100644
--- a/hipamd/src/hip_mempool.cpp
+++ b/hipamd/src/hip_mempool.cpp
@@ -59,7 +59,9 @@ hipError_t hipDeviceSetMemPool(int device, hipMemPool_t mem_pool) {
   if ((mem_pool == nullptr) || (device >= g_devices.size())) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-
+  if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
+    HIP_RETURN(hipErrorStreamCaptureUnsupported);
+  }
   auto poolDevice = reinterpret_cast<hip::MemoryPool*>(mem_pool)->Device();
   if (poolDevice->deviceId() != device) {
     HIP_RETURN(hipErrorInvalidDevice);
@@ -324,7 +326,9 @@ hipError_t hipMemPoolCreate(hipMemPool_t* mem_pool, const hipMemPoolProps* pool_
   if (IS_WINDOWS && pool_props->handleTypes == hipMemHandleTypePosixFileDescriptor) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-
+  if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
+    HIP_RETURN(hipErrorStreamCaptureUnsupported);
+  }
   auto device = g_devices[pool_props->location.id];
   auto pool = new hip::MemoryPool(device, pool_props);
   if (pool == nullptr) {
@@ -340,6 +344,9 @@ hipError_t hipMemPoolDestroy(hipMemPool_t mem_pool) {
   if (mem_pool == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+  if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
+    HIP_RETURN(hipErrorStreamCaptureUnsupported);
+  }
   hip::MemoryPool* hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
 
   if (!IsMemPoolValid(hip_mem_pool)) {