pthread-affinity-test.cpp

#include <pthread.h>
#include <sched.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/syscall.h>

#include <iostream>
#include <bitset>

#define gettid() syscall(SYS_gettid)

uint8_t get_affinity_mask()
{
    cpu_set_t cpu_set;
    if (sched_getaffinity(gettid(), sizeof(cpu_set_t), &cpu_set) == -1)
    {
        std::cout << "Failed: sched_getaffinity()" << std::endl;
        return -1;
    }

    uint8_t affinityMask = 0;
    for (int i=0; i<8; i++)
    {
        if (CPU_ISSET(i, &cpu_set))
        {
            affinityMask |= (1 << i);
        }
    }

    return affinityMask;
}

void *thread_func(void *args)
{
    uint8_t affinityMask = get_affinity_mask();

    std::bitset<8> affinityMaskBin(affinityMask);
    std::cout << "CPU affinity mask is: " << affinityMaskBin << std::endl;

    pthread_exit(0);
}

int main(int argc, char* argv[])
{
    void* res;

    pthread_t thread1;
    pthread_t thread2;

    pthread_attr_t thread1_attr;
    pthread_attr_t thread2_attr;

    uint8_t affinityMask = get_affinity_mask();

    std::bitset<8> affinityMaskBin(affinityMask);
    std::cout << "CPU affinity mask of main process is: " << affinityMaskBin << std::endl;

    if ((pthread_attr_init(&thread1_attr) != 0) ||
        (pthread_attr_init(&thread2_attr) != 0))
    {
        std::cout << "Failed: pthread_attr_init()" << std::endl;
    }

    cpu_set_t cpu_set;
    CPU_ZERO(&cpu_set);
    CPU_SET(1, &cpu_set);
    CPU_SET(3, &cpu_set);

    if (pthread_attr_setaffinity_np(&thread1_attr, sizeof(cpu_set_t), &cpu_set) != 0)
    {
        std::cout << "Failed: pthread_attr_setaffinity_np()" << std::endl;
    }

    CPU_ZERO(&cpu_set);
    CPU_SET(1, &cpu_set);
    CPU_SET(2, &cpu_set);
    CPU_SET(3, &cpu_set);

    if (pthread_attr_setaffinity_np(&thread2_attr, sizeof(cpu_set_t), &cpu_set) != 0)
    {
        std::cout << "Failed: pthread_attr_setaffinity_np()" << std::endl;
    }

    if (pthread_create(&thread1, &thread1_attr, thread_func, NULL) != 0)
    {
        std::cout << "Failed: pthread_create()" << std::endl;
    }

    if (pthread_join(thread1, &res) != 0)
    {
        std::cout << "Failed: pthread_join()" << std::endl;
    }

    if (pthread_create(&thread2, &thread2_attr, thread_func, NULL) != 0)
    {
        std::cout << "Failed: pthread_create()" << std::endl;
    }

    if (pthread_join(thread2, &res) != 0)
    {
        std::cout << "Failed: pthread_join()" << std::endl;
    }

    return 0;
}