| 25 | // ----------------------------------------------------------------------------------------------- |
| 26 | |
| 27 | ThreadAffinity::ThreadAffinity(int maxNumThreadsPerCore, int verbose) |
| 28 | : Verbose(verbose) |
| 29 | { |
| 30 | HMODULE hLib = GetModuleHandle(TEXT("kernel32")); |
| 31 | pGetLogicalProcessorInformationEx = |
| 32 | (GetLogicalProcessorInformationExFunc)GetProcAddress(hLib, "GetLogicalProcessorInformationEx"); |
| 33 | pSetThreadGroupAffinity = |
| 34 | (SetThreadGroupAffinityFunc)GetProcAddress(hLib, "SetThreadGroupAffinity"); |
| 35 | if (!pGetLogicalProcessorInformationEx || !pSetThreadGroupAffinity) |
| 36 | return; |
| 37 | |
| 38 | // Get logical processor information |
| 39 | PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX buffer = nullptr; |
| 40 | DWORD bufferSize = 0; |
| 41 | |
| 42 | // First call the function with an empty buffer to get the required buffer size |
| 43 | BOOL result = pGetLogicalProcessorInformationEx(RelationProcessorCore, buffer, &bufferSize); |
| 44 | if (result || GetLastError() != ERROR_INSUFFICIENT_BUFFER) |
| 45 | { |
| 46 | printWarning("GetLogicalProcessorInformationEx failed"); |
| 47 | return; |
| 48 | } |
| 49 | |
| 50 | // Allocate the buffer |
| 51 | buffer = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)malloc(bufferSize); |
| 52 | if (!buffer) |
| 53 | { |
| 54 | printWarning("SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX allocation failed"); |
| 55 | return; |
| 56 | } |
| 57 | |
| 58 | // Call again the function but now with the properly sized buffer |
| 59 | result = pGetLogicalProcessorInformationEx(RelationProcessorCore, buffer, &bufferSize); |
| 60 | if (!result) |
| 61 | { |
| 62 | printWarning("GetLogicalProcessorInformationEx failed"); |
| 63 | free(buffer); |
| 64 | return; |
| 65 | } |
| 66 | |
| 67 | // Iterate over the logical processor information structures |
| 68 | // There should be one structure for each physical core |
| 69 | char* ptr = (char*)buffer; |
| 70 | while (ptr < (char*)buffer + bufferSize) |
| 71 | { |
| 72 | PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX item = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)ptr; |
| 73 | if (item->Relationship == RelationProcessorCore && item->Processor.GroupCount > 0) |
| 74 | { |
| 75 | // Iterate over the groups |
| 76 | int numThreadsPerCore = 0; |
| 77 | for (int group = 0; group < item->Processor.GroupCount && |
| 78 | numThreadsPerCore < maxNumThreadsPerCore; ++group) |
| 79 | { |
| 80 | GROUP_AFFINITY coreAffinity = item->Processor.GroupMask[group]; |
| 81 | while (coreAffinity.Mask != 0 && numThreadsPerCore < maxNumThreadsPerCore) |
| 82 | { |
| 83 | // Extract the next set bit/thread from the mask |
| 84 | GROUP_AFFINITY threadAffinity = coreAffinity; |