mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-17 21:10:43 +00:00
113 lines
3.2 KiB
C++
113 lines
3.2 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#pragma once
|
|
|
|
#include <mutex>
|
|
#include <winrt/windows.media.h>
|
|
#include "VideoFrameToTensorConverter.h"
|
|
#include "TensorToVideoFrameConverter.h"
|
|
|
|
namespace _winml {
|
|
|
|
// Forward Declare
|
|
class ConverterResourceStore;
|
|
|
|
struct ConverterResourceDescription {
|
|
DWORD pixel_format;
|
|
int width;
|
|
int height;
|
|
LUID luid;
|
|
|
|
bool operator==(_In_ ConverterResourceDescription& desc) {
|
|
// Converter resources DON'T match if
|
|
// 1) the resources have different dimensions
|
|
// 2) the resources are on different devices
|
|
// 3) the resources have different pixel formats
|
|
if (desc.width != width || desc.height != height || desc.luid.HighPart != luid.HighPart ||
|
|
desc.luid.LowPart != luid.LowPart || desc.pixel_format != pixel_format) {
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
};
|
|
|
|
class ConverterResources : public std::enable_shared_from_this<ConverterResources> {
|
|
using Pool = std::weak_ptr<ConverterResourceStore>;
|
|
|
|
public:
|
|
template <typename... TArgs>
|
|
static std::shared_ptr<ConverterResources> Create(Pool pool, ConverterResourceDescription& descriptor) {
|
|
return std::make_shared<ConverterResources>(pool, descriptor);
|
|
}
|
|
|
|
ConverterResources(Pool& pool, ConverterResourceDescription& descriptor);
|
|
|
|
void ReturnToCache();
|
|
|
|
public:
|
|
ConverterResourceDescription Descriptor;
|
|
|
|
std::unique_ptr<_winml::VideoFrameToTensorConverter> Tensorizer;
|
|
std::unique_ptr<_winml::TensorToVideoFrameConverter> Detensorizer;
|
|
|
|
private:
|
|
Pool m_pool;
|
|
};
|
|
|
|
// This class retains tensorization and detensorization
|
|
// resources in a store, and evicts the oldest resource object
|
|
// when the size of the pool is maxed out. Objects in the pool
|
|
// can be reused for caching purposes to enhance performance during
|
|
// tensorization.
|
|
class ConverterResourceStore : public std::enable_shared_from_this<ConverterResourceStore> {
|
|
struct PoolObject {
|
|
std::shared_ptr<ConverterResources> Resource;
|
|
uint64_t StoreId;
|
|
};
|
|
|
|
public:
|
|
template <typename... TArgs>
|
|
static std::shared_ptr<ConverterResourceStore> Create(TArgs&&... args) {
|
|
return std::make_shared<ConverterResourceStore>(std::forward<TArgs>(args)...);
|
|
}
|
|
|
|
ConverterResourceStore(size_t nCacheSize);
|
|
|
|
std::shared_ptr<ConverterResources> Fetch(ConverterResourceDescription& descriptor);
|
|
void Store(std::shared_ptr<ConverterResources> object);
|
|
|
|
private:
|
|
std::shared_ptr<ConverterResources> FetchAndRemoveObject(ConverterResourceDescription& desc);
|
|
void EvictOldestPoolObject();
|
|
|
|
private:
|
|
std::vector<PoolObject> m_objects;
|
|
size_t m_cacheSize;
|
|
std::mutex m_mutex;
|
|
uint64_t storeId = 0;
|
|
};
|
|
|
|
class PoolObjectWrapper {
|
|
public:
|
|
template <typename... TArgs>
|
|
static std::shared_ptr<PoolObjectWrapper> Create(TArgs&&... args) {
|
|
return std::make_shared<PoolObjectWrapper>(std::forward<TArgs>(args)...);
|
|
}
|
|
|
|
explicit PoolObjectWrapper(std::shared_ptr<ConverterResources>&& resources) : m_resources(resources) {}
|
|
|
|
~PoolObjectWrapper() {
|
|
if (m_resources) {
|
|
m_resources->ReturnToCache();
|
|
}
|
|
}
|
|
|
|
std::shared_ptr<ConverterResources> Get() { return m_resources; }
|
|
|
|
private:
|
|
std::shared_ptr<ConverterResources> m_resources;
|
|
};
|
|
|
|
} // namespace _winml
|