You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
641 lines
21 KiB
641 lines
21 KiB
#pragma once
|
|
|
|
#include <iostream>
|
|
|
|
#include <unordered_map>
|
|
#include <shared_mutex>
|
|
#include <mutex>
|
|
#include <memory>
|
|
|
|
#include <sched.h>
|
|
#include <numa.h>
|
|
#include <numaif.h>
|
|
|
|
#include <dml/dml.hpp>
|
|
|
|
namespace dml {
|
|
inline const std::string StatusCodeToString(const dml::status_code code) {
|
|
switch (code) {
|
|
case dml::status_code::ok:
|
|
return "ok";
|
|
case dml::status_code::false_predicate:
|
|
return "false predicate";
|
|
case dml::status_code::partial_completion:
|
|
return "partial completion";
|
|
case dml::status_code::nullptr_error:
|
|
return "nullptr error";
|
|
case dml::status_code::bad_size:
|
|
return "bad size";
|
|
case dml::status_code::bad_length:
|
|
return "bad length";
|
|
case dml::status_code::inconsistent_size:
|
|
return "inconsistent size";
|
|
case dml::status_code::dualcast_bad_padding:
|
|
return "dualcast bad padding";
|
|
case dml::status_code::bad_alignment:
|
|
return "bad alignment";
|
|
case dml::status_code::buffers_overlapping:
|
|
return "buffers overlapping";
|
|
case dml::status_code::delta_delta_empty:
|
|
return "delta delta empty";
|
|
case dml::status_code::batch_overflow:
|
|
return "batch overflow";
|
|
case dml::status_code::execution_failed:
|
|
return "execution failed";
|
|
case dml::status_code::unsupported_operation:
|
|
return "unsupported operation";
|
|
case dml::status_code::queue_busy:
|
|
return "queue busy";
|
|
case dml::status_code::error:
|
|
return "unknown error";
|
|
case dml::status_code::config_error:
|
|
return "config error";
|
|
default:
|
|
return "unhandled error";
|
|
}
|
|
}
|
|
}
|
|
|
|
namespace dsacache {
|
|
class Cache;
|
|
|
|
// cache data holds all required information on
|
|
// one cache entry and will both be stored
|
|
// internally by the cache and handed out
|
|
// as copies to the user
|
|
// this class uses its object lifetime and
|
|
// a global reference counter to allow
|
|
// thread-safe copies and resource management
|
|
|
|
class CacheData {
|
|
public:
|
|
using dml_handler = dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>>;
|
|
|
|
private:
|
|
// data source and size of the block
|
|
uint8_t* src_;
|
|
size_t size_;
|
|
|
|
// global reference counting object
|
|
std::atomic<int32_t>* active_;
|
|
|
|
// global cache-location pointer
|
|
std::atomic<uint8_t*>* cache_;
|
|
|
|
// object-local incomplete cache location pointer
|
|
// which is only available in the first instance
|
|
uint8_t* incomplete_cache_;
|
|
|
|
// dml handler vector pointer which is only
|
|
// available in the first instance
|
|
std::unique_ptr<std::vector<dml_handler>> handlers_;
|
|
|
|
// deallocates the global cache-location
|
|
// and invalidates it
|
|
void Deallocate();
|
|
|
|
// checks whether there are at least two
|
|
// valid references to this object which
|
|
// is done as the cache always has one
|
|
// internally to any living instance
|
|
bool Active() const;
|
|
|
|
friend Cache;
|
|
public:
|
|
CacheData(uint8_t* data, const size_t size);
|
|
CacheData(const CacheData& other);
|
|
~CacheData();
|
|
|
|
// waits on completion of caching operations
|
|
// for this task and is safe to be called in
|
|
// any state of the object
|
|
void WaitOnCompletion();
|
|
|
|
// returns the cache data location for this
|
|
// instance which is valid as long as the
|
|
// instance is alive - !!! this may also
|
|
// yield a nullptr !!!
|
|
uint8_t* GetDataLocation() const;
|
|
};
|
|
|
|
// cache class will handle access to data through the cache
|
|
// by managing the cache through work submission, it sticks
|
|
// to user-defined caching and copy policies, is thread
|
|
// safe after initialization and returns copies of
|
|
// cache data class to the user
|
|
|
|
class Cache {
|
|
public:
|
|
// cache policy is defined as a type here to allow flexible usage of the cacher
|
|
// given a numa destination node (where the data will be needed), the numa source
|
|
// node (current location of the data) and the data size, this function should
|
|
// return optimal cache placement
|
|
// dst node and returned value can differ if the system, for example, has HBM
|
|
// attached accessible directly to node n under a different node id m
|
|
typedef int (CachePolicy)(const int numa_dst_node, const int numa_src_node, const size_t data_size);
|
|
|
|
// copy policy specifies the copy-executing nodes for a given task
|
|
// which allows flexibility in assignment for optimizing raw throughput
|
|
// or choosing a conservative usage policy
|
|
typedef std::vector<int> (CopyPolicy)(const int numa_dst_node, const int numa_src_node);
|
|
|
|
private:
|
|
// mutex for accessing the cache state map
|
|
|
|
std::shared_mutex cache_mutex_;
|
|
|
|
// map from [dst-numa-node,map2]
|
|
// map2 from [data-ptr,cache-structure]
|
|
|
|
std::unordered_map<uint8_t, std::unordered_map<uint8_t*, CacheData>> cache_state_;
|
|
|
|
CachePolicy* cache_policy_function_ = nullptr;
|
|
CopyPolicy* copy_policy_function_ = nullptr;
|
|
|
|
// function used to submit a copy task on a specific node to the dml
|
|
// engine on that node - will change the current threads node assignment
|
|
// to achieve this so take care to restore this
|
|
dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>> ExecuteCopy(
|
|
const uint8_t* src, uint8_t* dst, const size_t size, const int node
|
|
) const;
|
|
|
|
// allocates the required memory on the destination node
|
|
// and then submits task to the dml library for processing
|
|
// and attaches the handlers to the cache data structure
|
|
void SubmitTask(CacheData* task, const int dst_node, const int src_node);
|
|
|
|
// querries the policy functions for the given data and size
|
|
// to obtain destination cache node, also returns the datas
|
|
// source node for further usage
|
|
// output may depend on the calling threads node assignment
|
|
// as this is set as the "optimal placement" node
|
|
void GetCacheNode(uint8_t* src, const size_t size, int* OUT_DST_NODE, int* OUT_SRC_NODE) const;
|
|
|
|
// allocates memory of size "size" on the numa node "node"
|
|
// and returns nullptr if this is not possible, also may
|
|
// try to flush the cache of the requested node to
|
|
// alleviate encountered shortage
|
|
uint8_t* AllocOnNode(const size_t size, const int node);
|
|
|
|
// checks whether the cache contains an entry for
|
|
// the given data in the given memory node and
|
|
// returns it, otherwise returns nullptr
|
|
std::unique_ptr<CacheData> GetFromCache(uint8_t* src, const size_t size, const int dst_node);
|
|
|
|
public:
|
|
// initializes the cache with the two policy functions
|
|
// only after this is it safe to use in a threaded environment
|
|
void Init(CachePolicy* cache_policy_function, CopyPolicy* copy_policy_function);
|
|
|
|
// function to perform data access through the cache
|
|
std::unique_ptr<CacheData> Access(uint8_t* data, const size_t size);
|
|
|
|
// flushes the cache of inactive entries
|
|
// if node is -1 then the whole cache is
|
|
// checked and otherwise the specified
|
|
// node - no checks on node validity
|
|
void Flush(const int node = -1);
|
|
|
|
// forces out all entries from the
|
|
// cache and therefore will also "forget"
|
|
// still-in-use entries, these will still
|
|
// be properly deleted, but the cache
|
|
// will be fresh - use for testing
|
|
void Clear();
|
|
};
|
|
}
|
|
|
|
inline void dsacache::Cache::Clear() {
|
|
std::unique_lock<std::shared_mutex> lock(cache_mutex_);
|
|
|
|
cache_state_.clear();
|
|
|
|
Init(cache_policy_function_, copy_policy_function_);
|
|
}
|
|
|
|
inline void dsacache::Cache::Init(CachePolicy* cache_policy_function, CopyPolicy* copy_policy_function) {
|
|
cache_policy_function_ = cache_policy_function;
|
|
copy_policy_function_ = copy_policy_function;
|
|
|
|
// initialize numa library
|
|
|
|
numa_available();
|
|
|
|
// obtain all available nodes
|
|
// and those we may allocate
|
|
// memory on
|
|
|
|
const int nodes_max = numa_num_configured_nodes();
|
|
const bitmask* valid_nodes = numa_get_mems_allowed();
|
|
|
|
// prepare the cache state with entries
|
|
// for all given nodes
|
|
|
|
for (int node = 0; node < nodes_max; node++) {
|
|
if (numa_bitmask_isbitset(valid_nodes, node)) {
|
|
cache_state_.insert({node,{}});
|
|
}
|
|
}
|
|
}
|
|
|
|
inline std::unique_ptr<dsacache::CacheData> dsacache::Cache::Access(uint8_t* data, const size_t size) {
|
|
// get destination numa node for the cache
|
|
|
|
int dst_node = -1;
|
|
int src_node = -1;
|
|
|
|
GetCacheNode(data, size, &dst_node, &src_node);
|
|
|
|
// TODO: at this point it could be beneficial to check whether
|
|
// TODO: the given destination node is present as an entry
|
|
// TODO: in the cache state to see if it is valid
|
|
|
|
// check whether the data is already cached
|
|
|
|
std::unique_ptr<CacheData> task = GetFromCache(data, size, dst_node);
|
|
|
|
if (task != nullptr) {
|
|
return std::move(task);
|
|
}
|
|
|
|
// at this point the requested data is not present in cache
|
|
// and we create a caching task for it
|
|
|
|
task = std::make_unique<CacheData>(data, size);
|
|
|
|
{
|
|
std::unique_lock<std::shared_mutex> lock(cache_mutex_);
|
|
|
|
const auto state = cache_state_[dst_node].emplace(task->src_, *task);
|
|
|
|
// if state.second is false then no insertion took place
|
|
// which means that concurrently whith this thread
|
|
// some other thread must have accessed the same
|
|
// resource in which case we return the other
|
|
// threads data cache structure
|
|
|
|
if (!state.second) {
|
|
std::cout << "[!] Found another cache instance for 0x" << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
|
|
return std::move(std::make_unique<CacheData>(state.first->second));
|
|
}
|
|
}
|
|
|
|
SubmitTask(task.get(), dst_node, src_node);
|
|
|
|
return std::move(task);
|
|
}
|
|
|
|
inline uint8_t* dsacache::Cache::AllocOnNode(const size_t size, const int node) {
|
|
// allocate data on this node and flush the unused parts of the
|
|
// cache if the operation fails and retry once
|
|
// TODO: smarter flush strategy could keep some stuff cached
|
|
|
|
// check currently free memory to see if the data fits
|
|
|
|
long long int free_space = 0;
|
|
numa_node_size64(node, &free_space);
|
|
|
|
if (free_space < size) {
|
|
std::cout << "[!] Memory shortage when allocating " << size << "B on node " << node << std::endl;
|
|
|
|
// dst node lacks memory space so we flush the cache for this
|
|
// node hoping to free enough currently unused entries to make
|
|
// the second allocation attempt successful
|
|
|
|
Flush(node);
|
|
|
|
// re-test by getting the free space and checking again
|
|
|
|
numa_node_size64(node, &free_space);
|
|
|
|
if (free_space < size) {
|
|
std::cout << "[x] Memory shortage after flush when allocating " << size << "B on node " << node << std::endl;
|
|
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
uint8_t* dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(size, node));
|
|
|
|
if (dst == nullptr) {
|
|
std::cout << "[x] Allocation try failed for " << size << "B on node " << node << std::endl;
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
return dst;
|
|
}
|
|
|
|
inline void dsacache::Cache::SubmitTask(CacheData* task, const int dst_node, const int src_node) {
|
|
uint8_t* dst = AllocOnNode(task->size_, dst_node);
|
|
|
|
if (dst == nullptr) {
|
|
std::cout << "[x] Allocation failed so we can not cache" << std::endl;
|
|
return;
|
|
}
|
|
|
|
task->incomplete_cache_ = dst;
|
|
|
|
// querry copy policy function for the nodes to use for the copy
|
|
|
|
const std::vector<int> executing_nodes = copy_policy_function_(dst_node, src_node);
|
|
const size_t task_count = executing_nodes.size();
|
|
|
|
// each task will copy one fair part of the total size
|
|
// and in case the total size is not a factor of the
|
|
// given task count the last node must copy the remainder
|
|
|
|
const size_t size = task->size_ / task_count;
|
|
const size_t last_size = size + task->size_ % task_count;
|
|
|
|
// save the current numa node mask to restore later
|
|
// as executing the copy task will place this thread
|
|
// on a different node
|
|
|
|
bitmask* nodemask = numa_get_run_node_mask();
|
|
|
|
for (uint32_t i = 0; i < task_count; i++) {
|
|
const size_t local_size = i + 1 == task_count ? size : last_size;
|
|
const size_t local_offset = i * size;
|
|
const uint8_t* local_src = task->src_ + local_offset;
|
|
uint8_t* local_dst = dst + local_offset;
|
|
|
|
task->handlers_->emplace_back(ExecuteCopy(local_src, local_dst, local_size, executing_nodes[i]));
|
|
}
|
|
|
|
// restore the previous nodemask
|
|
|
|
numa_run_on_node_mask(nodemask);
|
|
numa_free_nodemask(nodemask);
|
|
}
|
|
|
|
inline dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>> dsacache::Cache::ExecuteCopy(
|
|
const uint8_t* src, uint8_t* dst, const size_t size, const int node
|
|
) const {
|
|
dml::const_data_view srcv = dml::make_view(src, size);
|
|
dml::data_view dstv = dml::make_view(dst, size);
|
|
|
|
numa_run_on_node(node);
|
|
|
|
return dml::submit<dml::automatic>(dml::mem_copy.block_on_fault(), srcv, dstv);
|
|
}
|
|
|
|
inline void dsacache::Cache::GetCacheNode(uint8_t* src, const size_t size, int* OUT_DST_NODE, int* OUT_SRC_NODE) const {
|
|
// obtain numa node of current thread to determine where the data is needed
|
|
|
|
const int current_cpu = sched_getcpu();
|
|
const int current_node = numa_node_of_cpu(current_cpu);
|
|
|
|
// obtain node that the given data pointer is allocated on
|
|
|
|
*OUT_SRC_NODE = -1;
|
|
get_mempolicy(OUT_SRC_NODE, NULL, 0, (void*)src, MPOL_F_NODE | MPOL_F_ADDR);
|
|
|
|
// querry cache policy function for the destination numa node
|
|
|
|
*OUT_DST_NODE = cache_policy_function_(current_node, *OUT_SRC_NODE, size);
|
|
}
|
|
|
|
inline void dsacache::Cache::Flush(const int node) {
|
|
// this lambda is used because below we have two code paths that
|
|
// flush nodes, either one single or all successively
|
|
|
|
const auto FlushNode = [](std::unordered_map<uint8_t*,CacheData>& map) {
|
|
// begin at the front of the map
|
|
|
|
auto it = map.begin();
|
|
|
|
// loop until we reach the end of the map
|
|
|
|
while (it != map.end()) {
|
|
// if the iterator points to an inactive element
|
|
// then we may erase it
|
|
|
|
if (it->second.Active() == false) {
|
|
// erase the iterator from the map
|
|
|
|
map.erase(it);
|
|
|
|
// as the erasure invalidated out iterator
|
|
// we must start at the beginning again
|
|
|
|
it = map.begin();
|
|
}
|
|
else {
|
|
// if element is active just move over to the next one
|
|
|
|
it++;
|
|
}
|
|
}
|
|
};
|
|
|
|
{
|
|
// we require exclusive lock as we modify the cache state
|
|
|
|
std::unique_lock<std::shared_mutex> lock(cache_mutex_);
|
|
|
|
// node == -1 means that cache on all nodes should be flushed
|
|
|
|
if (node == -1) {
|
|
for (auto& nc : cache_state_) {
|
|
FlushNode(nc.second);
|
|
}
|
|
}
|
|
else {
|
|
FlushNode(cache_state_[node]);
|
|
}
|
|
}
|
|
}
|
|
|
|
inline std::unique_ptr<dsacache::CacheData> dsacache::Cache::GetFromCache(uint8_t* src, const size_t size, const int dst_node) {
|
|
// the best situation is if this data is already cached
|
|
// which we check in an unnamed block in which the cache
|
|
// is locked for reading to prevent another thread
|
|
// from marking the element we may find as unused and
|
|
// clearing it
|
|
|
|
// lock the cache state in shared-mode because we read
|
|
|
|
std::shared_lock<std::shared_mutex> lock(cache_mutex_);
|
|
|
|
// search for the data in our cache state structure at the given node
|
|
|
|
const auto search = cache_state_[dst_node].find(src);
|
|
|
|
// if the data is in our structure we continue
|
|
|
|
if (search != cache_state_[dst_node].end()) {
|
|
|
|
// now check whether the sizes match
|
|
|
|
if (search->second.size_ >= size) {
|
|
// return a unique copy of the entry which uses the object
|
|
// lifetime and destructor to safely handle deallocation
|
|
|
|
return std::move(std::make_unique<CacheData>(search->second));
|
|
}
|
|
else {
|
|
// if the sizes missmatch then we clear the current entry from cache
|
|
// which will cause its deletion only after the last possible outside
|
|
// reference is also destroyed
|
|
|
|
cache_state_[dst_node].erase(search);
|
|
}
|
|
}
|
|
|
|
return nullptr;
|
|
}
|
|
|
|
inline dsacache::CacheData::CacheData(uint8_t* data, const size_t size) {
|
|
src_ = data;
|
|
size_ = size;
|
|
active_ = new std::atomic<int32_t>(1);
|
|
cache_ = new std::atomic<uint8_t*>();
|
|
incomplete_cache_ = nullptr;
|
|
handlers_ = std::make_unique<std::vector<dml_handler>>();
|
|
}
|
|
|
|
inline dsacache::CacheData::CacheData(const dsacache::CacheData& other) {
|
|
// we copy the ptr to the global atomic reference counter
|
|
// and increase the amount of active references
|
|
|
|
active_ = other.active_;
|
|
const int current_active = active_->fetch_add(1);
|
|
|
|
// source and size will be copied too
|
|
// as well as the reference to the global
|
|
// atomic cache pointer
|
|
|
|
src_ = other.src_;
|
|
size_ = other.size_;
|
|
cache_ = other.cache_;
|
|
|
|
// incomplete cache and handlers will not
|
|
// be copied because only the first instance
|
|
// will wait on the completion of handlers
|
|
|
|
incomplete_cache_ = nullptr;
|
|
handlers_ = nullptr;
|
|
}
|
|
|
|
inline dsacache::CacheData::~CacheData() {
|
|
// if this is the first instance of this cache structure
|
|
// and it has not been waited on and is now being destroyed
|
|
// we must wait on completion here to ensure the cache
|
|
// remains in a valid state
|
|
|
|
if (handlers_ != nullptr) {
|
|
WaitOnCompletion();
|
|
}
|
|
|
|
// due to fetch_sub returning the preivously held value
|
|
// we must subtract one locally to get the current value
|
|
|
|
const int32_t v = active_->fetch_sub(1) - 1;
|
|
|
|
// if the returned value is zero or lower
|
|
// then we must execute proper deletion
|
|
// as this was the last reference
|
|
|
|
if (v <= 0) {
|
|
Deallocate();
|
|
|
|
delete active_;
|
|
delete cache_;
|
|
}
|
|
}
|
|
|
|
inline void dsacache::CacheData::Deallocate() {
|
|
// although deallocate should only be called from
|
|
// a safe context to do so, it can not hurt to
|
|
// defensively perform the operation atomically
|
|
|
|
uint8_t* cache_local = cache_->exchange(nullptr);
|
|
if (cache_local != nullptr) numa_free(cache_local, size_);
|
|
}
|
|
|
|
inline uint8_t* dsacache::CacheData::GetDataLocation() const {
|
|
return cache_->load();
|
|
}
|
|
|
|
inline bool dsacache::CacheData::Active() const {
|
|
// this entry is active if more than one
|
|
// reference exists to it, as the Cache
|
|
// will always keep one internally until
|
|
// the entry is cleared from cache
|
|
|
|
return active_->load() > 1;
|
|
}
|
|
|
|
inline void dsacache::CacheData::WaitOnCompletion() {
|
|
// the cache data entry can be in two states
|
|
// either it is the original one which has not
|
|
// been waited for in which case the handlers
|
|
// are non-null or it is not
|
|
|
|
if (handlers_ == nullptr) {
|
|
// when no handlers are attached to this cache entry we wait on a
|
|
// value change for the cache structure from nullptr to non-null
|
|
// which will either go through immediately if the cache is valid
|
|
// already or wait until the handler-owning thread notifies us
|
|
|
|
cache_->wait(nullptr);
|
|
}
|
|
else {
|
|
// when the handlers are non-null there are some DSA task handlers
|
|
// available on which we must wait here
|
|
|
|
// abort is set if any operation encountered an error
|
|
|
|
bool abort = false;
|
|
|
|
for (auto& handler : *handlers_) {
|
|
auto result = handler.get();
|
|
|
|
if (result.status != dml::status_code::ok) {
|
|
std::cerr << "[x] Encountered bad status code for operation: " << dml::StatusCodeToString(result.status) << std::endl;
|
|
|
|
// if one of the copy tasks failed we abort the whole task
|
|
// after all operations are completed on it
|
|
|
|
abort = true;
|
|
}
|
|
}
|
|
|
|
// the handlers are cleared after all have completed
|
|
|
|
handlers_ = nullptr;
|
|
|
|
// now we act depending on whether an abort has been
|
|
// called for which signals operation incomplete
|
|
|
|
if (abort) {
|
|
// store nullptr in the cache location
|
|
|
|
cache_->store(nullptr);
|
|
|
|
// then free the now incomplete cache
|
|
|
|
// TODO: it would be possible to salvage the
|
|
// TODO: operation at this point but this
|
|
// TODO: is quite complicated so we just abort
|
|
|
|
numa_free(incomplete_cache_, size_);
|
|
}
|
|
else {
|
|
// incomplete cache is now safe to use and therefore we
|
|
// swap it with the global cache state of this entry
|
|
// and notify potentially waiting threads
|
|
|
|
cache_->store(incomplete_cache_);
|
|
}
|
|
|
|
// as a last step all waiting threads must
|
|
// be notified (copies of this will wait on value
|
|
// change of the cache) and the incomplete cache
|
|
// is cleared to nullptr as it is not incomplete
|
|
|
|
cache_->notify_all();
|
|
incomplete_cache_ = nullptr;
|
|
}
|
|
}
|