36 lines
679 B
C++
36 lines
679 B
C++
#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);
|
|
};
|