Move l0 fetcher completely into device class

Author: akroviakov

omniscidb/L0Mgr/L0Mgr.cpp

@@ -139,8 +139,7 @@ void* allocate_device_mem(const size_t num_bytes, L0Device& device) {
   return mem;
 }
 
-L0DataFetcher::L0DataFetcher(const L0Driver& driver, ze_device_handle_t device)
-    : device_(device), driver_(driver) {
+L0Device::L0DataFetcher::L0DataFetcher(L0Device& device) : my_device_(device) {
   ze_command_queue_desc_t command_queue_fetch_desc = {
       ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC,
       nullptr,
@@ -149,77 +148,87 @@ L0DataFetcher::L0DataFetcher(const L0Driver& driver, ze_device_handle_t device)
       0,
       ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS,
       ZE_COMMAND_QUEUE_PRIORITY_NORMAL};
-  L0_SAFE_CALL(zeCommandQueueCreate(
-      driver.ctx(), device_, &command_queue_fetch_desc, &queue_handle_));
-  current_cl_bytes = {{}, 0};
-  L0_SAFE_CALL(
-      zeCommandListCreate(driver.ctx(), device_, &cl_desc, &current_cl_bytes.first));
-}
-
-L0DataFetcher::~L0DataFetcher() {
+  L0_SAFE_CALL(zeCommandQueueCreate(my_device_.driver_.ctx(),
+                                    my_device_.device_,
+                                    &command_queue_fetch_desc,
+                                    &queue_handle_));
+  cur_cl_bytes_ = {{}, 0};
+  L0_SAFE_CALL(zeCommandListCreate(my_device_.driver_.ctx(),
+                                   my_device_.device_,
+                                   &cl_desc_,
+                                   &cur_cl_bytes_.cl_handle_));
+}
+
+L0Device::L0DataFetcher::~L0DataFetcher() {
   zeCommandQueueDestroy(queue_handle_);
-  zeCommandListDestroy(current_cl_bytes.first);
-  for (auto& dead_handle : graveyard) {
+  zeCommandListDestroy(cur_cl_bytes_.cl_handle_);
+  for (auto& dead_handle : graveyard_) {
     zeCommandListDestroy(dead_handle);
   }
-  for (auto& cl_handle : recycled) {
+  for (auto& cl_handle : recycled_) {
     zeCommandListDestroy(cl_handle);
   }
 }
 
-void L0DataFetcher::recycleGraveyard() {
-  while (recycled.size() < GRAVEYARD_LIMIT && graveyard.size()) {
-    recycled.push_back(graveyard.front());
-    graveyard.pop_front();
-    L0_SAFE_CALL(zeCommandListReset(recycled.back()));
+void L0Device::L0DataFetcher::recycleGraveyard() {
+  while (recycled_.size() < GRAVEYARD_LIMIT && graveyard_.size()) {
+    recycled_.push_back(graveyard_.front());
+    graveyard_.pop_front();
+    L0_SAFE_CALL(zeCommandListReset(recycled_.back()));
   }
-  for (auto& dead_handle : graveyard) {
-    L0_SAFE_CALL(zeCommandListDestroy(recycled.back()));
+  for (auto& dead_handle : graveyard_) {
+    L0_SAFE_CALL(zeCommandListDestroy(dead_handle));
+  }
+  graveyard_.clear();
+}
+
+void L0Device::L0DataFetcher::setCLRecycledOrNew() {
+  cur_cl_bytes_ = {{}, 0};
+  if (recycled_.size()) {
+    cur_cl_bytes_.cl_handle_ = recycled_.front();
+    recycled_.pop_front();
+  } else {
+    L0_SAFE_CALL(zeCommandListCreate(my_device_.driver_.ctx(),
+                                     my_device_.device_,
+                                     &cl_desc_,
+                                     &cur_cl_bytes_.cl_handle_));
   }
-  graveyard.clear();
 }
 
-void L0DataFetcher::appendCopyCommand(void* dst,
-                                      const void* src,
-                                      const size_t num_bytes) {
-  std::unique_lock<std::mutex> cl_lock(current_cl_lock);
+void L0Device::L0DataFetcher::appendCopyCommand(void* dst,
+                                                const void* src,
+                                                const size_t num_bytes) {
+  std::unique_lock<std::mutex> cl_lock(cur_cl_lock_);
   L0_SAFE_CALL(zeCommandListAppendMemoryCopy(
-      current_cl_bytes.first, dst, src, num_bytes, nullptr, 0, nullptr));
-  current_cl_bytes.second += num_bytes;
-  if (current_cl_bytes.second >= 128 * 1024 * 1024) {
-    ze_command_list_handle_t cl_h_copy = current_cl_bytes.first;
-    graveyard.push_back(current_cl_bytes.first);
-    current_cl_bytes = {{}, 0};
-    if (recycled.size()) {
-      current_cl_bytes.first = recycled.front();
-      recycled.pop_front();
-    } else {
-      L0_SAFE_CALL(
-          zeCommandListCreate(driver_.ctx(), device_, &cl_desc, &current_cl_bytes.first));
-    }
+      cur_cl_bytes_.cl_handle_, dst, src, num_bytes, nullptr, 0, nullptr));
+  cur_cl_bytes_.bytes_ += num_bytes;
+  if (cur_cl_bytes_.bytes_ >= CL_BYTES_LIMIT) {
+    ze_command_list_handle_t cl_h_copy = cur_cl_bytes_.cl_handle_;
+    graveyard_.push_back(cur_cl_bytes_.cl_handle_);
+    setCLRecycledOrNew();
     cl_lock.unlock();
     L0_SAFE_CALL(zeCommandListClose(cl_h_copy));
     L0_SAFE_CALL(
         zeCommandQueueExecuteCommandLists(queue_handle_, 1, &cl_h_copy, nullptr));
   }
 }
 
-void L0DataFetcher::sync() {
-  if (current_cl_bytes.second) {
-    L0_SAFE_CALL(zeCommandListClose(current_cl_bytes.first));
+void L0Device::L0DataFetcher::sync() {
+  if (cur_cl_bytes_.bytes_) {
+    L0_SAFE_CALL(zeCommandListClose(cur_cl_bytes_.cl_handle_));
     L0_SAFE_CALL(zeCommandQueueExecuteCommandLists(
-        queue_handle_, 1, &current_cl_bytes.first, nullptr));
+        queue_handle_, 1, &cur_cl_bytes_.cl_handle_, nullptr));
   }
   L0_SAFE_CALL(
       zeCommandQueueSynchronize(queue_handle_, std::numeric_limits<uint32_t>::max()));
-  L0_SAFE_CALL(zeCommandListReset(current_cl_bytes.first));
-  if (graveyard.size() > GRAVEYARD_LIMIT) {
+  L0_SAFE_CALL(zeCommandListReset(cur_cl_bytes_.cl_handle_));
+  if (graveyard_.size() > GRAVEYARD_LIMIT) {
     recycleGraveyard();
   }
 }
 
 L0Device::L0Device(const L0Driver& driver, ze_device_handle_t device)
-    : device_(device), driver_(driver), data_fetcher(driver, device) {
+    : device_(device), driver_(driver), data_fetcher_(*this) {
   ze_command_queue_handle_t queue_handle;
   ze_command_queue_desc_t command_queue_desc = {ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC,
                                                 nullptr,
@@ -271,6 +280,14 @@ unsigned L0Device::maxSharedLocalMemory() const {
   return compute_props_.maxSharedLocalMemory;
 }
 
+void L0Device::transferToDevice(void* dst, const void* src, const size_t num_bytes) {
+  data_fetcher_.appendCopyCommand(dst, src, num_bytes);
+}
+
+void L0Device::syncDataTransfers() {
+  data_fetcher_.sync();
+}
+
 L0CommandQueue::L0CommandQueue(ze_command_queue_handle_t handle) : handle_(handle) {}
 
 ze_command_queue_handle_t L0CommandQueue::handle() const {
@@ -420,7 +437,7 @@ void L0Manager::copyHostToDeviceAsync(int8_t* device_ptr,
   CHECK_LT(device_num, drivers_[0]->devices().size());
 
   auto& device = drivers()[0]->devices()[device_num];
-  device->data_fetcher.appendCopyCommand(device_ptr, host_ptr, num_bytes);
+  device->transferToDevice(device_ptr, host_ptr, num_bytes);
 }
 
 void L0Manager::copyHostToDeviceAsyncIfPossible(int8_t* device_ptr,
@@ -438,7 +455,7 @@ void L0Manager::synchronizeDeviceDataStream(const int device_num) {
   CHECK_GE(device_num, 0);
   CHECK_LT(device_num, drivers_[0]->devices().size());
   auto& device = drivers()[0]->devices()[device_num];
-  device->data_fetcher.sync();
+  device->syncDataTransfers();
 }
 
 void L0Manager::copyDeviceToHost(int8_t* host_ptr,

omniscidb/L0Mgr/L0Mgr.h

@@ -58,51 +58,65 @@ class L0Kernel;
 class L0CommandList;
 class L0CommandQueue;
 
-class L0DataFetcher {
+class L0Device {
+ protected:
+  const L0Driver& driver_;
+  std::shared_ptr<L0CommandQueue> command_queue_;
 #ifdef HAVE_L0
-  static constexpr uint16_t GRAVEYARD_LIMIT{500};
   ze_device_handle_t device_;
-  ze_command_queue_handle_t queue_handle_;
-  std::pair<ze_command_list_handle_t, uint64_t> current_cl_bytes;
-  std::list<ze_command_list_handle_t> graveyard;
-  std::list<ze_command_list_handle_t> recycled;
-  ze_command_list_desc_t cl_desc = {ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC,
-                                    nullptr,
-                                    0,
-                                    ZE_COMMAND_LIST_FLAG_MAXIMIZE_THROUGHPUT};
-  std::mutex current_cl_lock;
+  ze_device_properties_t props_;
+  ze_device_compute_properties_t compute_props_;
 #endif
-  const L0Driver& driver_;
-  void recycleGraveyard();
-
- public:
-  void appendCopyCommand(void* dst, const void* src, const size_t num_bytes);
-  void sync();
-
+ private:
+  /*
+    This component for data fetching to L0 devices is used for
+    more efficient asynchronous data transfers. It allows to amortize
+    the costs of data transfers which is especially useful in case of
+    many relatively small data transfers.
+  */
+  class L0DataFetcher {
 #ifdef HAVE_L0
-  L0DataFetcher(const L0Driver& driver, ze_device_handle_t device);
-  ~L0DataFetcher();
-#else
-  L0DataFetcher() = default;
+    static constexpr size_t GRAVEYARD_LIMIT{500};
+    static constexpr size_t CL_BYTES_LIMIT{128 * 1024 * 1024};
+    static constexpr ze_command_list_desc_t cl_desc_ = {
+        ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC,
+        nullptr,
+        0,
+        ZE_COMMAND_LIST_FLAG_MAXIMIZE_THROUGHPUT};
+    struct CLBytesTracker {
+      ze_command_list_handle_t cl_handle_;
+      uint64_t bytes_;
+    };
+    CLBytesTracker cur_cl_bytes_;
+    ze_command_queue_handle_t queue_handle_;
+    std::list<ze_command_list_handle_t> graveyard_;
+    std::list<ze_command_list_handle_t> recycled_;
+    std::mutex cur_cl_lock_;
 #endif
-};
+    L0Device& my_device_;
+    void recycleGraveyard();
+    void setCLRecycledOrNew();
+
+   public:
+    void appendCopyCommand(void* dst, const void* src, const size_t num_bytes);
+    void sync();
 
-class L0Device {
- private:
 #ifdef HAVE_L0
-  ze_device_handle_t device_;
-  ze_device_properties_t props_;
-  ze_device_compute_properties_t compute_props_;
+    L0DataFetcher(L0Device& device);
+    ~L0DataFetcher();
+#else
+    L0DataFetcher() = default;
 #endif
-
-  const L0Driver& driver_;
-  std::shared_ptr<L0CommandQueue> command_queue_;
+  };
+  L0DataFetcher data_fetcher_;
 
  public:
-  L0DataFetcher data_fetcher;
   std::shared_ptr<L0CommandQueue> command_queue() const;
   std::unique_ptr<L0CommandList> create_command_list() const;
 
+  void transferToDevice(void* dst, const void* src, const size_t num_bytes);
+  void syncDataTransfers();
+
   std::shared_ptr<L0Module> create_module(uint8_t* code,
                                           size_t len,
                                           bool log = false) const;