xc-llm-ascend/csrc/vnpu_offload/shm_worker.h

#pragma once

#include <vector>
#include <atomic>
#include <thread>

#include "shm_helper.h"


class ShmWorker {
 public:
  ShmWorker();
  ~ShmWorker();

  bool register_worker(int32_t tgid, int gpu_id, uint64_t *out_shareable_handle,
                       uint64_t *out_vmem_size);

  bool try_lock_gpu(bool &out_self_hold);
  bool lock_gpu(bool &out_self_hold);
  void unlock_gpu();

 private:
  int32_t tgid;
  int gpu_id;
  ShmHelper *shm_helper;
  std::thread heart_beat_thread;
  std::atomic<bool> stop_heart_beat;

  // request
  uint64_t make_request(uint32_t type, uint64_t parameter);
  int register_worker_shm();

  // heart beat
  void heart_beat_loop(int slot);
};