106 lines
2.8 KiB
C++
106 lines
2.8 KiB
C++
#include "Core/Threads.h"
|
|
#include <iostream>
|
|
#include <vector>
|
|
#include <cassert>
|
|
|
|
#ifdef __linux__
|
|
#include <pthread.h>
|
|
#include <sched.h>
|
|
#endif
|
|
|
|
using namespace uLib;
|
|
|
|
std::vector<int> GetAvailableCpus() {
|
|
std::vector<int> available;
|
|
#ifdef __linux__
|
|
cpu_set_t cpuset;
|
|
CPU_ZERO(&cpuset);
|
|
if (sched_getaffinity(0, sizeof(cpu_set_t), &cpuset) == 0) {
|
|
for (int i = 0; i < CPU_SETSIZE; ++i) {
|
|
if (CPU_ISSET(i, &cpuset)) {
|
|
available.push_back(i);
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
return available;
|
|
}
|
|
|
|
class TestThread : public Thread {
|
|
public:
|
|
void Run() override {
|
|
Thread::Sleep(200);
|
|
}
|
|
};
|
|
|
|
void TestThreadAffinity() {
|
|
std::cout << "Testing Thread Affinity..." << std::endl;
|
|
#ifdef __linux__
|
|
auto available = GetAvailableCpus();
|
|
if (available.empty()) {
|
|
std::cout << " No CPUs available for affinity test, skipping." << std::endl;
|
|
return;
|
|
}
|
|
int target_cpu = available[0];
|
|
std::cout << " Using CPU " << target_cpu << std::endl;
|
|
|
|
TestThread t;
|
|
t.Start();
|
|
t.SetAffinity(target_cpu);
|
|
|
|
cpu_set_t cpuset;
|
|
CPU_ZERO(&cpuset);
|
|
int s = pthread_getaffinity_np(t.GetNativeHandle(), sizeof(cpu_set_t), &cpuset);
|
|
if (s != 0) {
|
|
std::cerr << "Error: pthread_getaffinity_np failed with code " << s << std::endl;
|
|
assert(false);
|
|
}
|
|
assert(CPU_ISSET(target_cpu, &cpuset));
|
|
|
|
t.Join();
|
|
std::cout << " Passed (Thread bound to CPU 0)." << std::endl;
|
|
#else
|
|
std::cout << " Affinity not supported on this OS, skipping." << std::endl;
|
|
#endif
|
|
}
|
|
|
|
void TestTeamAffinity() {
|
|
std::cout << "Testing Team Affinity..." << std::endl;
|
|
#ifdef __linux__
|
|
auto available = GetAvailableCpus();
|
|
if (available.size() < 2) {
|
|
std::cout << " Not enough CPUs available for Team affinity test, skipping." << std::endl;
|
|
return;
|
|
}
|
|
std::vector<int> cpus = {available[0], available[1]};
|
|
std::cout << " Using CPUs " << cpus[0] << ", " << cpus[1] << std::endl;
|
|
|
|
Team team(2);
|
|
team.SetAffinity(cpus);
|
|
|
|
// We check affinity inside a parallel region
|
|
#pragma omp parallel
|
|
{
|
|
cpu_set_t cpuset;
|
|
CPU_ZERO(&cpuset);
|
|
pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
|
|
int tid = Thread::GetThreadNum();
|
|
int expected_cpu = cpus[tid % cpus.size()];
|
|
assert(CPU_ISSET(expected_cpu, &cpuset));
|
|
}
|
|
std::cout << " Passed (Team threads bound correctly)." << std::endl;
|
|
#else
|
|
std::cout << " Affinity not supported on this OS, skipping." << std::endl;
|
|
#endif
|
|
}
|
|
|
|
// Helper to get native handle if needed (oops, I forgot to add it to Thread class)
|
|
// I'll add GetNativeHandle() to Thread class in Threads.h
|
|
|
|
int main() {
|
|
TestThreadAffinity();
|
|
TestTeamAffinity();
|
|
std::cout << "All Affinity tests finished!" << std::endl;
|
|
return 0;
|
|
}
|