temp: add xPUBackend refine source tree

TODO:
complete xpubackend registry
add more interface for xpu chips
export some header files

refine

Author: lipracer

CMakeLists.txt

@@ -67,3 +67,4 @@ target_link_libraries(af
         ${TORCH_PYTHON_LIBRARY})
 
 add_subdirectory(tests)
+add_subdirectory(plugins/klx_backend)

fserver/csrc/public.hpp

@@ -1,4 +1,7 @@
 /* Copyright (c) 2025, StepFun Authors. All rights reserved. */
+
+#include <dlfcn.h>
+
 #include <execinfo.h>
 #include <stdio.h>
 #include <signal.h>
@@ -43,6 +46,7 @@ uint64_t handler_counter_ = 0;
 std::unordered_map<uint64_t, AFTensorMeta> meta_map_;
 std::vector<std::deque<ServerDataBatch>> q_;
 std::atomic<uint64_t> q_signal_;
+static void* gPluginHandle = nullptr;
 
 void RequestHandler(const AFTensorMeta& req_meta, AFTensorServer* server) {
   std::vector<torch::Tensor> tensors;
@@ -165,7 +169,12 @@ void barrier(bool include_server, bool include_worker, bool instrance_barrier=tr
 }
 
 
-void init() {
+void init(const std::string& plugin) {
+  if (!plugin.empty()) {
+    gPluginHandle = dlopen(plugin.c_str(), RTLD_NOW);
+    PS_CHECK(gPluginHandle)
+        << "can't load plugin:" << plugin << ": " << dlerror();
+  }
 
   std::string role_str = ps::GetEnv("DMLC_ROLE", "server");
   int offset = 0;
@@ -204,14 +213,17 @@ void stop() {
     ps::Postoffice::GetWorker(gpu_)->DoBarrier(0,
         ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler, true);
   } else if (role_ == Node::SERVER) {
-    ps::Postoffice::GetServer(gpu_)->DoBarrier(0,
-        ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler, true);
+    ps::Postoffice::GetServer(gpu_)->DoBarrier(
+        0, ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler, true);
   } else {
-    ps::Postoffice::Get()->DoBarrier(0,
-        ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler, true);
+    ps::Postoffice::Get()->DoBarrier(
+        0, ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler, true);
   }
 
   ps::Finalize(0, role_, true);
+  if (gPluginHandle) {
+    dlclose(gPluginHandle);
+  }
 }
 
 std::vector<int> get_all_handlers(int handler) {
@@ -238,7 +250,8 @@ uint64_t get_nanosecond() {
 
 
 void pybind_public(py::module &m){
-  m.def("init", &init, py::call_guard<py::gil_scoped_release>());
+  m.def("init", &init, py::arg("plugin") = "",
+        py::call_guard<py::gil_scoped_release>());
   m.def("stop", &stop, py::call_guard<py::gil_scoped_release>());
 
   m.def("register_recv_buffer",

include/dmlc/backend_registry.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "base.h"
+#include "ps/backend.h"
+
+namespace dmlc {
+
+template <typename T>
+struct STEPMESH_API backend_registry {
+  backend_registry(const std::string& name) {
+    ps::Backend::RegisterLazy(name, []() { return new T(); });
+  }
+};
+
+}  // namespace dmlc

include/dmlc/base.h

@@ -190,4 +190,6 @@ inline const char *BeginPtr(const std::string &str) {
 #define alignof __alignof
 #endif
 
+#define STEPMESH_API __attribute__((__visibility__("default")))
+
 #endif  // DMLC_BASE_H_

include/dmlc/logging.h

@@ -71,25 +71,35 @@ inline void InitLogging(const char *argv0) {
   // DO NOTHING
 }
 
+constexpr const char *getFileName(const char *path) {
+  auto last = path + strlen(path);
+  while (*last != '/') {
+    --last;
+  }
+  return ++last;
+}
+
 // Always-on checking
-#define PS_CHECK(x)                                                      \
-  if (!(x))                                                           \
-  dmlc::LogMessageFatal(__FILE__, __LINE__).stream() << "Check "      \
-                                                        "failed: " #x \
-                                                     << ' '
+#define PS_CHECK(x)                                                     \
+  if (!(x))                                                             \
+  dmlc::LogMessageFatal(dmlc::getFileName(__FILE__), __LINE__).stream() \
+      << "Check "                                                       \
+         "failed: " #x                                                  \
+      << ' '
 #define PS_CHECK_LT(x, y) PS_CHECK((x) < (y))
 #define PS_CHECK_GT(x, y) PS_CHECK((x) > (y))
 #define PS_CHECK_LE(x, y) PS_CHECK((x) <= (y))
 #define PS_CHECK_GE(x, y) PS_CHECK((x) >= (y))
 #define PS_CHECK_EQ(x, y) PS_CHECK((x) == (y))
 #define PS_CHECK_NE(x, y) PS_CHECK((x) != (y))
-#define PS_CHECK_NOTNULL(x)                                                                 \
-  ((x) == NULL                                                                           \
-   ? dmlc::LogMessageFatal(__FILE__, __LINE__).stream() << "Check  notnull: " #x << ' ', \
+#define PS_CHECK_NOTNULL(x)                                                \
+  ((x) == NULL                                                             \
+   ? dmlc::LogMessageFatal(dmlc::getFileName(__FILE__), __LINE__).stream() \
+         << "Check  notnull: " #x << ' ',                                  \
    (x) : (x))  // NOLINT(*)
 // Debug-only checking.
 #ifdef NDEBUG
-/* 
+/*
 #define DPS_CHECK(x) \
   while (false) PS_CHECK(x)
 #define DPS_CHECK_LT(x, y) \
@@ -114,12 +124,12 @@ inline void InitLogging(const char *argv0) {
 #define DPS_CHECK_NE(x, y) PS_CHECK((x) != (y)) */
 #endif  // NDEBUG
 
-#define PS_LOG_API dmlc::LogMessage(__FILE__, __LINE__)
+#define PS_LOG_API dmlc::LogMessage(dmlc::getFileName(__FILE__), __LINE__)
 
 #define PS_LOG_IF(severity, condition) \
   !(condition) ? (void)0 : dmlc::LogMessageVoidify() & PS_LOG_API
 
-#define LOG_FATAL dmlc::LogMessageFatal(__FILE__, __LINE__)
+#define LOG_FATAL dmlc::LogMessageFatal(dmlc::getFileName(__FILE__), __LINE__)
 #define PS_LOG_FATAL LOG_FATAL.stream()
 #define LOG_QFATAL LOG_FATAL
 

include/ps/af_tensor_app.h

@@ -20,7 +20,7 @@
 
 #include "ps/base.h"
 #include "ps/hash_table8.hpp"
-#include "ps/internal/backend.h"
+#include "ps/backend.h"
 #include "ps/internal/utils.h"
 #include "ps/kv_app.h"
 
@@ -236,15 +236,16 @@ class AFTensorWorker {
   void ZPush_(int ts, const SArray<Key>& keys, const at::Tensor& tensor,
               int cmd = 0) {
     SArray<char> val;
-    val.reset(reinterpret_cast<char*>(tensor.data_ptr()),
+    void* mappedPtr = Backend::Get()->GetAccessibleAddr(tensor);
+    val.reset(reinterpret_cast<char*>(mappedPtr),
               tensor.numel() * tensor.itemsize(), [tensor](void*) {});
 
     Message msg;
     msg.meta.request = true;
     msg.meta.head = cmd;
     msg.meta.push = true;
     msg.meta.timestamp = ts;
-    msg.meta.addr = reinterpret_cast<uint64_t>(tensor.data_ptr());
+    msg.meta.addr = reinterpret_cast<uint64_t>(mappedPtr);
     msg.meta.val_len = tensor.numel() * tensor.itemsize();
     PS_VLOG(2) << "ZPush_ addr: 0x" << std::hex << msg.meta.addr << std::dec
                << " val_len: " << msg.meta.val_len;
@@ -284,13 +285,14 @@ class AFTensorWorker {
 
       *key.data() = pull_tensors[i * pull_batch_size + index].key;
 
-      val.reset(reinterpret_cast<char*>(tensor.data_ptr()),
+      void* mappedPtr = Backend::Get()->GetAccessibleAddr(tensor);
+      val.reset(reinterpret_cast<char*>(mappedPtr),
                 tensor.numel() * tensor.itemsize(), [tensor](void*) {});
 
       msg.meta.request = true;
       msg.meta.head = cmd;
       msg.meta.push = false;
-      msg.meta.addr = reinterpret_cast<uint64_t>(tensor.data_ptr());
+      msg.meta.addr = reinterpret_cast<uint64_t>(mappedPtr);
       msg.meta.val_len = tensor.numel() * tensor.itemsize();
       msg.meta.key = key[0];
       msg.meta.is_tensor = 1;
@@ -483,7 +485,8 @@ class AFTensorServer {
         res.keys = key;
 
         SArray<char> tensor_val;
-        tensor_val.reset(reinterpret_cast<char*>(tensors[0].val.data_ptr()),
+        tensor_val.reset(reinterpret_cast<char*>(
+                             Backend::Get()->GetAccessibleAddr(tensors[0].val)),
                          tensors[0].val.numel() * tensors[0].val.itemsize(),
                          [](void*) {});
         res.vals = tensor_val;
@@ -506,7 +509,8 @@ class AFTensorServer {
             rsp.kv_pair.keys = key;
 
             rsp.kv_pair.vals.reset(
-                reinterpret_cast<char*>(res_kv.val.data_ptr()),
+                reinterpret_cast<char*>(
+                    Backend::Get()->GetAccessibleAddr(res_kv.val)),
                 res_kv.val.numel() * res_kv.val.itemsize(), [](void*) {});
 
             rsp.kv_meta = kv_meta;
@@ -558,7 +562,8 @@ class AFTensorServer {
     PS_CHECK_GT(worker_ranks.size(), 0) << "ranks or keys should not be empty";
     PS_CHECK_EQ(worker_ranks.size(), keys.size())
         << "rank list and key list have unequal size";
-    char* buffer_ptr = reinterpret_cast<char*>(tensor.data_ptr());
+    char* buffer_ptr =
+        reinterpret_cast<char*>(Backend::Get()->GetAccessibleAddr(tensor));
     uint64_t data_size = tensor.numel() * tensor.element_size();
     int chunk_size = data_size / worker_ranks.size();
     PS_CHECK_EQ(data_size % worker_ranks.size(), 0)
@@ -591,8 +596,14 @@ class AFTensorServer {
                          .dtype(at::ScalarType(req_meta.dtype))
                          .memory_format(at::MemoryFormat::Contiguous)
                          .device(Backend::Get()->GetDevice());
-      key_tensor.val =
-          at::from_blob(req_data.vals.data(), req_meta.shape, options);
+      key_tensor.val = at::from_blob(
+          Backend::Get()->GetDeviceAddrFromHostPtr(
+              req_data.vals.data(),
+              std::accumulate(std::begin(req_meta.shape),
+                              std::end(req_meta.shape),
+                              c10::elementSize(at::ScalarType(req_meta.dtype)),
+                              std::multiplies<uint64_t>())),
+          req_meta.shape, options);
     }
     key_tensor.key = req_data.keys[0];
     return key_tensor;

include/ps/internal/backend.h include/ps/backend.hinclude/ps/internal/backend.h renamed to include/ps/backend.h

@@ -15,9 +15,11 @@
 #include <string>
 #include <unordered_map>
 #include <utility>
+#include <functional>
 
 #include "dmlc/logging.h"
-#include "ps/internal/env.h"
+#include "ps/env.h"
+#include "base.h"
 
 namespace ps {
 
@@ -26,7 +28,7 @@ enum { BACKEND_OK = 0, BACKEND_FAILED = -1 };
 /**
  * \brief Abstract Backend Class
  */
-class Backend {
+class STEPMESH_API Backend {
  public:
   /**
    * \brief Set device index for current thread
@@ -88,6 +90,24 @@ class Backend {
    */
   virtual int SyncEvent(void* event) = 0;
 
+  /**
+   *\brief Get an address that is directly readable via the PCIe bus
+   * @param devicePtr device physical address
+   * @return an address that is directly readable via the PCIe bus
+   */
+  virtual void* GetAccessibleAddr(void* devicePtr, size_t size) {
+    return devicePtr;
+  }
+
+  virtual void* GetAccessibleAddr(const at::Tensor& tensor) {
+    return GetAccessibleAddr(tensor.data_ptr(),
+                             tensor.numel() * tensor.element_size());
+  }
+
+  virtual void* GetDeviceAddrFromHostPtr(void* hostPtr, size_t size) {
+    return hostPtr;
+  }
+
   /**
    * \brief Get the backend implementation
    * @return the backend implementation
@@ -98,21 +118,28 @@ class Backend {
     RegisterImpl(name, backend);
   }
 
+  static void RegisterLazy(const std::string& name,
+                           const std::function<Backend*(void)>& ctor);
  protected:
   Backend() = default;
 
  private:
   static std::mutex backends_mutex_;
   static std::unordered_map<std::string, Backend*> backends_;
+  static std::unordered_map<std::string, std::function<Backend*(void)>> backend_ctors_;
 
   static Backend* GetImpl() {
     static Backend* backend_impl = nullptr;
     if (backend_impl == nullptr) {
       std::unique_lock<std::mutex> lock(backends_mutex_);
       std::string backend_type = "GPU";
-      backend_type = Environment::Get()->find("STEPMESH_BAKCEND", backend_type);
-      PS_CHECK_NE(backends_.find(backend_type), backends_.end())
-          << "failed to get backend impl: " << backend_type;
+      backend_type = Environment::Get()->find("STEPMESH_BACKEND", backend_type);
+      if (backends_.find(backend_type) == backends_.end()) {
+        PS_CHECK_NE(backend_ctors_.find(backend_type), backend_ctors_.end())
+            << "failed to get backend impl: " << backend_type;
+        backends_[backend_type] = backend_ctors_[backend_type]();
+      }
+
       backend_impl = backends_[backend_type];
     }
     return backend_impl;

include/ps/internal/env.h include/ps/env.hinclude/ps/internal/env.h renamed to include/ps/env.h

@@ -6,6 +6,9 @@
  */
 #ifndef PS_INTERNAL_ASSIGN_OP_H_
 #define PS_INTERNAL_ASSIGN_OP_H_
+
+#include <cassert>
+
 #include "ps/internal/utils.h"
 namespace ps {
 
@@ -24,7 +27,7 @@ enum AssignOp {
  * \brief return an assignment function: right op= left
  */
 template <typename T>
-inline void AssignFunc(const T& lhs, AssignOp op, T* rhs) {
+inline void AssignFunc(const T& left, AssignOp op, T* right) {
   switch (op) {
     case ASSIGN:
       *right = left;
@@ -42,7 +45,7 @@ inline void AssignFunc(const T& lhs, AssignOp op, T* rhs) {
       *right /= left;
       break;
     default:
-      LOG(FATAL) << "use AssignOpInt..";
+      PS_LOG(FATAL) << "use AssignOpInt..";
   }
 }
 
@@ -51,7 +54,7 @@ inline void AssignFunc(const T& lhs, AssignOp op, T* rhs) {
  * works for integers
  */
 template <typename T>
-inline void AssignFuncInt(const T& lhs, AssignOp op, T* rhs) {
+inline void AssignFuncInt(const T& left, AssignOp op, T* right) {
   switch (op) {
     case ASSIGN:
       *right = left;

include/ps/internal/assign_op.h

@@ -7,7 +7,7 @@
 
 #include <memory>
 
-#include "ps/internal/backend.h"
+#include "ps/backend.h"
 
 namespace ps {