Refactor OVRTAllocator to return base pointer of remote tensor

This change allows for allocations made by the ov allocator to be
imported to other APIs that require base addresses to the original
device allocation.

Author: ericcraw

onnxruntime/core/providers/openvino/ov_allocator.cc

@@ -10,12 +10,6 @@ namespace onnxruntime {
 
 using namespace openvino_ep;
 
-constexpr size_t default_alignment = 4096;
-
-static inline size_t align_up(size_t size, size_t pow2_alignment) {
-  return (size + pow2_alignment - 1) & ~(pow2_alignment - 1);
-}
-
 OVRTAllocator::OVRTAllocator(ov::Core& core, OrtDevice::DeviceType device_type, OrtDevice::DeviceId device_id, const char* name) : IAllocator(OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, OrtDevice(device_type, OrtDevice::MemType::DEFAULT, device_id), device_id, OrtMemTypeCPUInput)), core_(core) {
   if (device_type == OrtDevice::NPU) {
     remote_ctx_ = core_.get_default_context("NPU").as<ov::intel_npu::level_zero::ZeroContext>();
@@ -26,25 +20,26 @@ OVRTAllocator::OVRTAllocator(ov::Core& core, OrtDevice::DeviceType device_type,
 
 void* OVRTAllocator::Alloc(size_t size) {
   try {
-    size_t alloc_size = align_up(size + sizeof(ov::Tensor*) + default_alignment, default_alignment);
     ov::Tensor* tensor = new ov::Tensor(remote_ctx_.create_host_tensor(ov::element::Type_t::u8,
-                                                                       {alloc_size}));
-    uintptr_t data_ptr = reinterpret_cast<uintptr_t>(tensor->data());
-
-    ov::Tensor** ptr = reinterpret_cast<ov::Tensor**>(align_up(data_ptr + sizeof(ov::Tensor*), default_alignment));
-    ptr[-1] = tensor;
-
-    return reinterpret_cast<void*>(ptr);
-
+                                                                       {size}));
+    std::unique_lock lock(mutex_);
+    allocated_.insert({tensor->data(), tensor});
+    return reinterpret_cast<void*>(tensor->data());
   } catch (const ov::Exception& e) {
     ORT_THROW(std::string("Alloc failed: ") + e.what());
   }
 }
 
 void OVRTAllocator::Free(void* p) {
   try {
-    ov::Tensor** ptr = reinterpret_cast<ov::Tensor**>(p);
-    delete ptr[-1];
+    std::unique_lock lock(mutex_);
+    auto it = allocated_.find(p);
+    if (it != allocated_.end()) {
+      ov::Tensor* tensor = it->second;
+      allocated_.erase(it);
+      lock.unlock();
+      delete tensor;
+    }
   } catch (const ov::Exception& e) {
     ORT_THROW(std::string("Free failed: ") + e.what());
   }

onnxruntime/core/providers/openvino/ov_allocator.h

@@ -3,9 +3,12 @@
 #ifdef USE_OVEP_NPU_MEMORY
 #pragma once
 
+#include <mutex>
+
 #include "core/common/inlined_containers.h"
 #include "core/framework/allocator.h"
 #include "openvino/runtime/remote_context.hpp"
+#include "core/common/inlined_containers.h"
 
 namespace onnxruntime {
 
@@ -18,6 +21,8 @@ class OVRTAllocator : public IAllocator {
  private:
   ov::Core& core_;
   ov::RemoteContext remote_ctx_;
+  InlinedHashMap<void*, ov::Tensor*> allocated_;
+  std::mutex mutex_;
 };
 
 }  // namespace onnxruntime