onnxruntime/winml/lib/Api.Image/inc/ConverterResourceStore.h
Justin Chu c203d89958
Update ruff and clang-format versions (#21479)
ruff -> 0.5.4
clang-format -> 18
2024-07-24 11:50:11 -07:00

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