From d13581efa7cb12b4d49d4054b2ec7b38fc1d9f72 Mon Sep 17 00:00:00 2001
From: Christophe Paquot <christophe.paquot@amd.com>
Date: Mon, 7 Jun 2021 14:14:08 -0700
Subject: [PATCH] SWDEV-276396 - Implement hipDeviceReset

Add a Purge function to MemObjMap

Change-Id: Iac51dfda9a7b7c45f2f4a0dc35f7a623121aba1a
---
 rocclr/device/device.cpp | 14 ++++++++++++++
 rocclr/device/device.hpp |  1 +
 2 files changed, 15 insertions(+)

diff --git a/rocclr/device/device.cpp b/rocclr/device/device.cpp
index bc951d9d61..0f163a267b 100644
--- a/rocclr/device/device.cpp
+++ b/rocclr/device/device.cpp
@@ -327,6 +327,20 @@ void MemObjMap::UpdateAccess(amd::Device *peerDev) {
   }
 }
 
+void MemObjMap::Purge(amd::Device* dev) {
+  assert(dev != nullptr);
+
+  amd::ScopedLock lock(AllocatedLock_);
+  for (auto it = MemObjMap_.cbegin() ; it != MemObjMap_.cend() ;) {
+    const std::vector<Device*>& devices = it->second->getContext().devices();
+    if (devices.size() == 1 && devices[0] == dev) {
+      it = MemObjMap_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 Device::BlitProgram::~BlitProgram() {
   if (program_ != nullptr) {
     program_->release();
diff --git a/rocclr/device/device.hpp b/rocclr/device/device.hpp
index 53c2e39fcb..71ebf5ea45 100644
--- a/rocclr/device/device.hpp
+++ b/rocclr/device/device.hpp
@@ -1238,6 +1238,7 @@ class MemObjMap : public AllStatic {
   static amd::Memory* FindMemObj(
       const void* k);  //!< find the mem object based on the input pointer
   static void UpdateAccess(amd::Device *peerDev);
+  static void Purge(amd::Device*dev); //!< Purge all the memories on the given device
  private:
   static std::map<uintptr_t, amd::Memory*>
       MemObjMap_;                      //!< the mem object<->hostptr information container