add kt-kernel

kt-kernel/cpu_backend/task_queue.cpp

@@ -0,0 +1,61 @@
+/**
+ * @Description :
+ * @Author    : chenht2022
+ * @Date     : 2024-07-17 12:25:51
+ * @Version   : 1.0.0
+ * @LastEditors : chenht2022
+ * @LastEditTime : 2024-10-09 11:08:10
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+#include "task_queue.h"
+
+#include <pthread.h>
+#include <sched.h>
+
+#include <chrono>
+#include <iostream>
+#include <thread>
+
+TaskQueue::TaskQueue() : done(false), pending(0) {
+  Node* dummy = new Node();
+  head.store(dummy, std::memory_order_relaxed);
+  tail.store(dummy, std::memory_order_relaxed);
+  workerThread = std::thread(&TaskQueue::worker, this);
+}
+
+TaskQueue::~TaskQueue() {
+  done.store(true, std::memory_order_release);
+  if (workerThread.joinable()) workerThread.join();
+
+  Node* node = head.load(std::memory_order_relaxed);
+  while (node) {
+    Node* next = node->next.load(std::memory_order_relaxed);
+    delete node;
+    node = next;
+  }
+}
+
+void TaskQueue::enqueue(std::function<void()> task) {
+  pending.fetch_add(1, std::memory_order_acq_rel);
+  Node* node = new Node(task);
+  Node* prev = tail.exchange(node, std::memory_order_acq_rel);
+  prev->next.store(node, std::memory_order_release);
+}
+
+void TaskQueue::sync(size_t n) { while (pending.load(std::memory_order_acquire) > n); }
+
+void TaskQueue::worker() {
+  Node* curr = head.load(std::memory_order_relaxed);
+  while (!done.load(std::memory_order_acquire)) {
+    Node* next = curr->next.load(std::memory_order_acquire);
+    if (next) {
+      if (next->task) {
+        next->task();
+      }
+      delete curr;
+      curr = next;
+      head.store(curr, std::memory_order_release);
+      pending.fetch_sub(1, std::memory_order_acq_rel);
+    }
+  }
+}