agent
15b18a8f1f
ThreadPinning reads the sibling groups from the kernel topology and pins the calling thread to a cpu set. It is stateless and per-thread: a caller reads the groups, picks the cpus for a thread and pins it, so any mapping of threads to cpus is expressible.
88 lines
2.4 KiB
C++
88 lines
2.4 KiB
C++
#include "ThreadPinning.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <fstream>
|
|
#include <set>
|
|
#include <sstream>
|
|
#include <string>
|
|
|
|
#if defined(__linux__) || defined(__ANDROID__)
|
|
#include <sched.h>
|
|
#define TP_THREAD_PINNING_SUPPORTED
|
|
#endif
|
|
|
|
namespace tp {
|
|
|
|
namespace {
|
|
std::vector<int> parseCpuList(const std::string& text) {
|
|
std::vector<int> cpus;
|
|
std::stringstream stream(text);
|
|
std::string range;
|
|
while (std::getline(stream, range, ',')) {
|
|
if (range.empty()) continue;
|
|
const size_t dash = range.find('-');
|
|
if (dash == std::string::npos) {
|
|
cpus.push_back(std::stoi(range));
|
|
continue;
|
|
}
|
|
const int first = std::stoi(range.substr(0, dash));
|
|
const std::string tail = range.substr(dash + 1);
|
|
const int last = tail.empty() ? first : std::stoi(tail);
|
|
for (int cpu = first; cpu <= last; cpu++)
|
|
cpus.push_back(cpu);
|
|
}
|
|
return cpus;
|
|
}
|
|
|
|
std::vector<int> siblingsOf(int cpu) {
|
|
std::ifstream file("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/thread_siblings_list");
|
|
std::string text;
|
|
if (!file || !std::getline(file, text)) return { cpu };
|
|
const std::vector<int> siblings = parseCpuList(text);
|
|
return siblings.empty() ? std::vector<int>{ cpu } : siblings;
|
|
}
|
|
}
|
|
|
|
std::vector<int> ThreadPinning::allowedCpus() {
|
|
#ifdef TP_THREAD_PINNING_SUPPORTED
|
|
cpu_set_t set;
|
|
CPU_ZERO(&set);
|
|
if (sched_getaffinity(0, sizeof(set), &set) != 0) return {};
|
|
std::vector<int> cpus;
|
|
for (int cpu = 0; cpu < CPU_SETSIZE; cpu++)
|
|
if (CPU_ISSET(cpu, &set)) cpus.push_back(cpu);
|
|
return cpus;
|
|
#else
|
|
return {};
|
|
#endif
|
|
}
|
|
|
|
std::vector<std::vector<int>> ThreadPinning::siblingGroups(const std::vector<int>& allowed) {
|
|
const std::vector<int> cpus = allowed.empty() ? allowedCpus() : allowed;
|
|
std::vector<std::vector<int>> groups;
|
|
std::set<int> seen;
|
|
for (int cpu : cpus) {
|
|
if (seen.count(cpu)) continue;
|
|
std::vector<int> group;
|
|
for (int sibling : siblingsOf(cpu)) {
|
|
if (!std::count(cpus.begin(), cpus.end(), sibling)) continue;
|
|
group.push_back(sibling);
|
|
seen.insert(sibling);
|
|
}
|
|
if (group.empty()) group.push_back(cpu);
|
|
groups.push_back(group);
|
|
}
|
|
return groups;
|
|
}
|
|
|
|
void ThreadPinning::pinCallingThread(const std::vector<int>& cpus) {
|
|
#ifdef TP_THREAD_PINNING_SUPPORTED
|
|
if (cpus.empty()) return;
|
|
cpu_set_t set;
|
|
CPU_ZERO(&set);
|
|
for (int cpu : cpus)
|
|
CPU_SET(cpu, &set);
|
|
sched_setaffinity(0, sizeof(set), &set);
|
|
#endif
|
|
}
|
|
}
|