This contains my bachelors thesis and associated tex files, code snippets and maybe more. Topic: Data Movement in Heterogeneous Memories with Intel Data Streaming Accelerator
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.

352 lines
13 KiB

  1. #pragma once
  2. #include <iostream>
  3. #include <unordered_map>
  4. #include <shared_mutex>
  5. #include <mutex>
  6. #include <memory>
  7. #include <sched.h>
  8. #include <numa.h>
  9. #include <numaif.h>
  10. #include <dml/dml.hpp>
  11. #include "cache-data.hpp"
  12. namespace dsacache {
  13. // cache class will handle access to data through the cache
  14. // by managing the cache through work submission, it sticks
  15. // to user-defined caching and copy policies, is thread
  16. // safe after initialization and returns copies of
  17. // cache data class to the user
  18. class Cache {
  19. public:
  20. // cache policy is defined as a type here to allow flexible usage of the cacher
  21. // given a numa destination node (where the data will be needed), the numa source
  22. // node (current location of the data) and the data size, this function should
  23. // return optimal cache placement
  24. // dst node and returned value can differ if the system, for example, has HBM
  25. // attached accessible directly to node n under a different node id m
  26. typedef int (CachePolicy)(const int numa_dst_node, const int numa_src_node, const size_t data_size);
  27. // copy policy specifies the copy-executing nodes for a given task
  28. // which allows flexibility in assignment for optimizing raw throughput
  29. // or choosing a conservative usage policy
  30. typedef std::vector<int> (CopyPolicy)(const int numa_dst_node, const int numa_src_node);
  31. private:
  32. // mutex for accessing the cache state map
  33. std::shared_mutex cache_mutex_;
  34. // map from [dst-numa-node,map2]
  35. // map2 from [data-ptr,cache-structure]
  36. std::unordered_map<uint8_t, std::unordered_map<uint8_t*, CacheData>> cache_state_;
  37. CachePolicy* cache_policy_function_ = nullptr;
  38. CopyPolicy* copy_policy_function_ = nullptr;
  39. // function used to submit a copy task on a specific node to the dml
  40. // engine on that node - will change the current threads node assignment
  41. // to achieve this so take care to restore this
  42. dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>> ExecuteCopy(
  43. const uint8_t* src, uint8_t* dst, const size_t size, const int node
  44. ) const;
  45. // allocates the required memory on the destination node
  46. // and then submits task to the dml library for processing
  47. // and attaches the handlers to the cache data structure
  48. void SubmitTask(CacheData* task, const int dst_node, const int src_node);
  49. // querries the policy functions for the given data and size
  50. // to obtain destination cache node, also returns the datas
  51. // source node for further usage
  52. // output may depend on the calling threads node assignment
  53. // as this is set as the "optimal placement" node
  54. void GetCacheNode(uint8_t* src, const size_t size, int* OUT_DST_NODE, int* OUT_SRC_NODE) const;
  55. // checks whether the cache contains an entry for
  56. // the given data in the given memory node and
  57. // returns it, otherwise returns nullptr
  58. std::unique_ptr<CacheData> GetFromCache(uint8_t* src, const size_t size, const int dst_node);
  59. public:
  60. // initializes the cache with the two policy functions
  61. // only after this is it safe to use in a threaded environment
  62. void Init(CachePolicy* cache_policy_function, CopyPolicy* copy_policy_function);
  63. // function to perform data access through the cache
  64. std::unique_ptr<CacheData> Access(uint8_t* data, const size_t size);
  65. // flushes the cache of inactive entries
  66. // if node is -1 then the whole cache is
  67. // checked and otherwise the specified
  68. // node - no checks on node validity
  69. void Flush(const int node = -1);
  70. };
  71. }
  72. inline void dsacache::Cache::Init(CachePolicy* cache_policy_function, CopyPolicy* copy_policy_function) {
  73. cache_policy_function_ = cache_policy_function;
  74. copy_policy_function_ = copy_policy_function;
  75. // initialize numa library
  76. numa_available();
  77. // obtain all available nodes
  78. // and those we may allocate
  79. // memory on
  80. const int nodes_max = numa_num_configured_nodes();
  81. const bitmask* valid_nodes = numa_get_mems_allowed();
  82. // prepare the cache state with entries
  83. // for all given nodes
  84. for (int node = 0; node < nodes_max; node++) {
  85. if (numa_bitmask_isbitset(valid_nodes, node)) {
  86. cache_state_.insert({node,{}});
  87. }
  88. }
  89. std::cout << "[-] Cache Initialized" << std::endl;
  90. }
  91. inline std::unique_ptr<dsacache::CacheData> dsacache::Cache::Access(uint8_t* data, const size_t size) {
  92. // get destination numa node for the cache
  93. int dst_node = -1;
  94. int src_node = -1;
  95. GetCacheNode(data, size, &dst_node, &src_node);
  96. // TODO: at this point it could be beneficial to check whether
  97. // TODO: the given destination node is present as an entry
  98. // TODO: in the cache state to see if it is valid
  99. // check whether the data is already cached
  100. std::unique_ptr<CacheData> task = GetFromCache(data, size, dst_node);
  101. if (task != nullptr) {
  102. return std::move(task);
  103. }
  104. // at this point the requested data is not present in cache
  105. // and we create a caching task for it
  106. task = std::make_unique<CacheData>(data, size);
  107. {
  108. std::unique_lock<std::shared_mutex> lock(cache_mutex_);
  109. const auto state = cache_state_[dst_node].emplace(task->src_, *task);
  110. // if state.second is false then no insertion took place
  111. // which means that concurrently whith this thread
  112. // some other thread must have accessed the same
  113. // resource in which case we return the other
  114. // threads data cache structure
  115. if (!state.second) {
  116. std::cout << "[!] Found another cache instance for 0x" << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
  117. return std::move(std::make_unique<CacheData>(state.first->second));
  118. }
  119. }
  120. SubmitTask(task.get(), dst_node, src_node);
  121. return std::move(task);
  122. }
  123. inline void dsacache::Cache::SubmitTask(CacheData* task, const int dst_node, const int src_node) {
  124. std::cout << "[+] Allocating " << task->size_ << "B on node " << dst_node << " for " << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
  125. // allocate data on this node and flush the unused parts of the
  126. // cache if the operation fails and retry once
  127. // TODO: smarter flush strategy could keep some stuff cached
  128. uint8_t* dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(task->size_, dst_node));
  129. if (dst == nullptr) {
  130. std::cout << "[!] First allocation try failed for " << task->size_ << "B on node " << dst_node << std::endl;
  131. // allocation on dst_node failed so we flush the cache for this
  132. // node hoping to free enough currently unused entries to make
  133. // the second allocation attempt successful
  134. Flush(dst_node);
  135. dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(task->size_, dst_node));
  136. if (dst == nullptr) {
  137. std::cerr << "[x] Second allocation try failed for " << task->size_ << "B on node " << dst_node << std::endl;
  138. return;
  139. }
  140. }
  141. task->incomplete_cache_ = dst;
  142. // querry copy policy function for the nodes to use for the copy
  143. const std::vector<int> executing_nodes = copy_policy_function_(dst_node, src_node);
  144. const size_t task_count = executing_nodes.size();
  145. // each task will copy one fair part of the total size
  146. // and in case the total size is not a factor of the
  147. // given task count the last node must copy the remainder
  148. const size_t size = task->size_ / task_count;
  149. const size_t last_size = size + task->size_ % task_count;
  150. std::cout << "[-] Splitting Copy into " << task_count << " tasks of " << size << "B 0x" << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
  151. // save the current numa node mask to restore later
  152. // as executing the copy task will place this thread
  153. // on a different node
  154. bitmask* nodemask = numa_get_run_node_mask();
  155. for (uint32_t i = 0; i < task_count; i++) {
  156. const size_t local_size = i + 1 == task_count ? size : last_size;
  157. const size_t local_offset = i * size;
  158. const uint8_t* local_src = task->src_ + local_offset;
  159. uint8_t* local_dst = dst + local_offset;
  160. task->handlers_->emplace_back(ExecuteCopy(local_src, local_dst, local_size, executing_nodes[i]));
  161. }
  162. // restore the previous nodemask
  163. numa_run_on_node_mask(nodemask);
  164. numa_free_nodemask(nodemask);
  165. }
  166. inline dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>> dsacache::Cache::ExecuteCopy(
  167. const uint8_t* src, uint8_t* dst, const size_t size, const int node
  168. ) const {
  169. dml::const_data_view srcv = dml::make_view(src, size);
  170. dml::data_view dstv = dml::make_view(dst, size);
  171. numa_run_on_node(node);
  172. return dml::submit<dml::automatic>(dml::mem_copy.block_on_fault(), srcv, dstv);
  173. }
  174. void dsacache::Cache::GetCacheNode(uint8_t* src, const size_t size, int* OUT_DST_NODE, int* OUT_SRC_NODE) const {
  175. // obtain numa node of current thread to determine where the data is needed
  176. const int current_cpu = sched_getcpu();
  177. const int current_node = numa_node_of_cpu(current_cpu);
  178. // obtain node that the given data pointer is allocated on
  179. *OUT_SRC_NODE = -1;
  180. get_mempolicy(OUT_SRC_NODE, NULL, 0, (void*)src, MPOL_F_NODE | MPOL_F_ADDR);
  181. // querry cache policy function for the destination numa node
  182. *OUT_DST_NODE = cache_policy_function_(current_node, *OUT_SRC_NODE, size);
  183. }
  184. inline void dsacache::Cache::Flush(const int node) {
  185. std::cout << "[-] Flushing Cache for " << (node == -1 ? "all nodes" : "node " + std::to_string(node)) << std::endl;
  186. // this lambda is used because below we have two code paths that
  187. // flush nodes, either one single or all successively
  188. const auto FlushNode = [](std::unordered_map<uint8_t*,CacheData>& map) {
  189. // begin at the front of the map
  190. auto it = map.begin();
  191. // loop until we reach the end of the map
  192. while (it != map.end()) {
  193. // if the iterator points to an inactive element
  194. // then we may erase it
  195. if (it->second.Active() == false) {
  196. // erase the iterator from the map
  197. map.erase(it);
  198. // as the erasure invalidated out iterator
  199. // we must start at the beginning again
  200. it = map.begin();
  201. }
  202. else {
  203. // if element is active just move over to the next one
  204. it++;
  205. }
  206. }
  207. };
  208. {
  209. // we require exclusive lock as we modify the cache state
  210. std::unique_lock<std::shared_mutex> lock(cache_mutex_);
  211. // node == -1 means that cache on all nodes should be flushed
  212. if (node == -1) {
  213. for (auto& nc : cache_state_) {
  214. FlushNode(nc.second);
  215. }
  216. }
  217. else {
  218. FlushNode(cache_state_[node]);
  219. }
  220. }
  221. }
  222. std::unique_ptr<dsacache::CacheData> dsacache::Cache::GetFromCache(uint8_t* src, const size_t size, const int dst_node) {
  223. // the best situation is if this data is already cached
  224. // which we check in an unnamed block in which the cache
  225. // is locked for reading to prevent another thread
  226. // from marking the element we may find as unused and
  227. // clearing it
  228. // lock the cache state in shared-mode because we read
  229. std::shared_lock<std::shared_mutex> lock(cache_mutex_);
  230. // search for the data in our cache state structure at the given node
  231. const auto search = cache_state_[dst_node].find(src);
  232. // if the data is in our structure we continue
  233. if (search != cache_state_[dst_node].end()) {
  234. // now check whether the sizes match
  235. // TODO: second.size_ >= size would also work
  236. if (search->second.size_ == size) {
  237. std::cout << "[+] Found Cached version for 0x" << std::hex << (uint64_t)src << std::dec << std::endl;
  238. // return a unique copy of the entry which uses the object
  239. // lifetime and destructor to safely handle deallocation
  240. return std::move(std::make_unique<CacheData>(search->second));
  241. }
  242. else {
  243. std::cout << "[!] Found Cached version with size missmatch for 0x" << std::hex << (uint64_t)src << std::dec << std::endl;
  244. // if the sizes missmatch then we clear the current entry from cache
  245. // which will cause its deletion only after the last possible outside
  246. // reference is also destroyed
  247. cache_state_[dst_node].erase(search);
  248. }
  249. }
  250. return nullptr;
  251. }