72 lines
1.8 KiB
C++
72 lines
1.8 KiB
C++
#pragma once
|
|
|
|
#include <cstdint>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "utils/simple_json.h"
|
|
|
|
#if defined(__linux__)
|
|
#include <pthread.h>
|
|
#include <sched.h>
|
|
#elif defined(_WIN32)
|
|
#ifndef NOMINMAX
|
|
#define NOMINMAX
|
|
#endif
|
|
#include <windows.h>
|
|
#endif
|
|
|
|
namespace rk3588 {
|
|
|
|
inline std::vector<int> ParseCpuAffinity(const SimpleJson& node_cfg) {
|
|
std::vector<int> out;
|
|
const SimpleJson* a = node_cfg.Find("cpu_affinity");
|
|
if (!a || !a->IsArray()) return out;
|
|
for (const auto& v : a->AsArray()) {
|
|
if (!v.IsNumber()) continue;
|
|
const int cpu = v.AsInt(-1);
|
|
if (cpu >= 0) out.push_back(cpu);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
inline bool SetCurrentThreadAffinity(const std::vector<int>& cpus, std::string& err) {
|
|
err.clear();
|
|
if (cpus.empty()) return true;
|
|
|
|
#if defined(__linux__)
|
|
cpu_set_t set;
|
|
CPU_ZERO(&set);
|
|
for (int cpu : cpus) {
|
|
if (cpu < 0 || cpu >= CPU_SETSIZE) continue;
|
|
CPU_SET(cpu, &set);
|
|
}
|
|
const int rc = pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
|
|
if (rc != 0) {
|
|
err = "pthread_setaffinity_np failed";
|
|
return false;
|
|
}
|
|
return true;
|
|
#elif defined(_WIN32)
|
|
// Windows uses a bitmask; best-effort for up to 64 CPUs in the current group.
|
|
DWORD_PTR mask = 0;
|
|
for (int cpu : cpus) {
|
|
if (cpu < 0 || cpu >= static_cast<int>(8 * sizeof(DWORD_PTR))) continue;
|
|
mask |= (static_cast<DWORD_PTR>(1) << cpu);
|
|
}
|
|
if (mask == 0) return true;
|
|
const DWORD_PTR prev = SetThreadAffinityMask(GetCurrentThread(), mask);
|
|
if (prev == 0) {
|
|
err = "SetThreadAffinityMask failed";
|
|
return false;
|
|
}
|
|
return true;
|
|
#else
|
|
(void)cpus;
|
|
err = "thread affinity not supported on this platform";
|
|
return false;
|
|
#endif
|
|
}
|
|
|
|
} // namespace rk3588
|