[AIMIGRAPHX-885] Add hip graph to MIGraphX EP

onnxruntime/core/providers/migraphx/migraphx_allocator.cc

@@ -22,18 +22,50 @@
 #endif
 }
 
+void MIGraphXAllocator::EnablePoolMode() {
+  std::lock_guard<std::mutex> lock(pool_mu_);
+  pool_enabled_ = true;
+}
+
 void* MIGraphXAllocator::Alloc(size_t size) {
   CheckDevice();
+  if (size == 0) return nullptr;
+
+  if (pool_enabled_) {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    auto it = free_list_.find(size);
+    if (it != free_list_.end() && !it->second.empty()) {
+      void* p = it->second.back();
+      it->second.pop_back();
+      return p;
+    }
+  }
+
   void* p = nullptr;
-  if (size > 0) {
-    HIP_CALL_THROW(hipMalloc((void**)&p, size));
+  HIP_CALL_THROW(hipMalloc((void**)&p, size));
+
+  if (pool_enabled_) {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    alloc_sizes_[p] = size;
   }
+
   return p;
 }
 
 void MIGraphXAllocator::Free(void* p) {
   CheckDevice();
-  (void)hipFree(p);  // do not throw error since it's OK for hipFree to fail during shutdown
+  if (!p) return;
+
+  if (pool_enabled_) {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    auto it = alloc_sizes_.find(p);
+    if (it != alloc_sizes_.end()) {
+      free_list_[it->second].push_back(p);
+      return;
+    }
+  }
+
+  (void)hipFree(p);
 }
 
 void* MIGraphXExternalAllocator::Alloc(size_t size) {

onnxruntime/core/providers/migraphx/migraphx_allocator.h

@@ -4,7 +4,9 @@
 #pragma once
 
 #include <mutex>
+#include <unordered_map>
 #include <unordered_set>
+#include <vector>
 #include "core/framework/allocator.h"
 
 namespace onnxruntime {
@@ -21,8 +23,20 @@ class MIGraphXAllocator : public IAllocator {
   virtual void* Alloc(size_t size) override;
   virtual void Free(void* p) override;
 
+  void EnablePoolMode();
+  bool IsPoolModeEnabled() const { return pool_enabled_; }
+
  private:
   void CheckDevice() const;
+
+  // When pool mode is enabled (for hipGraph), freed allocations are cached by
+  // size so that subsequent Alloc calls for the same size return the same
+  // device pointer.  This provides pointer stability required by hipGraph
+  // replay without the cost of intermediary buffer copies.
+  bool pool_enabled_ = false;
+  mutable std::mutex pool_mu_;
+  std::unordered_map<size_t, std::vector<void*>> free_list_;
+  std::unordered_map<void*, size_t> alloc_sizes_;
 };
 
 class MIGraphXExternalAllocator : public MIGraphXAllocator {