cpu_set.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <cstdint>
#include <thread>
#include <vector>
 
#include <dispenso/platform.h>
 
#if defined(__linux__) || defined(__FreeBSD__)
#define DISPENSO_CPUSET_LINUXY
#endif // linuxy
 
#if defined(DISPENSO_CPUSET_LINUXY)
#include <pthread.h>
#include <sched.h>
#elif defined(_WIN32)
#define DISPENSO_CPUSET_WINDOWS
#endif // supported os
 
// Note: Windows uses SetThreadGroupAffinity for binding and
// GetLogicalProcessorInformationEx for NUMA/cache topology, behind the same public API.
//
// macOS does not support explicit CPU pinning. bindCurrentThread() is a no-op and
// returns false. Topology queries return a single-node fallback.
//
// FreeBSD binding support is present (cpuset_setaffinity uses a similar interface to
// Linux), but topology detection requires test hardware and is deferred.
 
#ifndef DISPENSO_MAX_GROUP_SIZE
#define DISPENSO_MAX_GROUP_SIZE 16
#endif
 
namespace dispenso {
 
constexpr int32_t kDefaultMaxGroupSize = DISPENSO_MAX_GROUP_SIZE;
 
struct CacheGroup {
  std::vector<int32_t> cpus; 
  int32_t cacheId; 
};
 
// Forward declaration; full definition follows CpuSet (ThreadGroup has a CpuSet member).
struct ThreadGroup;
 
class CpuSet {
 public:
  DISPENSO_DLL_ACCESS CpuSet();
 
  DISPENSO_DLL_ACCESS void clear();
 
  DISPENSO_DLL_ACCESS void add(int32_t hardwareThread);
 
  DISPENSO_DLL_ACCESS void addRange(int32_t start, int32_t end);
 
  DISPENSO_DLL_ACCESS void remove(int32_t hardwareThread);
 
  DISPENSO_DLL_ACCESS void removeRange(int32_t start, int32_t end);
 
  DISPENSO_DLL_ACCESS bool contains(int32_t hardwareThread) const;
 
  DISPENSO_DLL_ACCESS int32_t count() const;
 
  DISPENSO_DLL_ACCESS bool bindCurrentThread() const;
 
  DISPENSO_DLL_ACCESS static int32_t totalNumaNodes();
 
  DISPENSO_DLL_ACCESS static int32_t currentHardwareThread();
 
  static int32_t currentHardwareThreadApprox() {
    static constexpr uint32_t kRefreshPeriod = 32;
    static DISPENSO_THREAD_LOCAL int32_t cachedCpu = -1;
    static DISPENSO_THREAD_LOCAL uint32_t counter = 0;
    if ((counter++ & (kRefreshPeriod - 1)) == 0) {
      cachedCpu = currentHardwareThread();
    }
    return cachedCpu;
  }
 
  DISPENSO_DLL_ACCESS static const CpuSet& node(int32_t numaNode);
 
  DISPENSO_DLL_ACCESS static const CpuSet& all();
 
  DISPENSO_DLL_ACCESS static int32_t availableCount();
 
  DISPENSO_DLL_ACCESS static const std::vector<CacheGroup>& l2CacheGroups();
 
  DISPENSO_DLL_ACCESS static const std::vector<CacheGroup>& l3CacheGroups();
 
  DISPENSO_DLL_ACCESS static std::vector<ThreadGroup> buildThreadGroups(
      int32_t maxGroupSize = kDefaultMaxGroupSize);
 
 private:
#if defined(DISPENSO_CPUSET_LINUXY)
  cpu_set_t set_;
#else
  // Portable bitset supporting up to kMaxCpus logical processors.
  // 1024 matches Linux cpu_set_t and covers all practical configurations.
  // Windows Server 2025 supports up to 12,288 (192 groups x 64) but
  // no known hardware approaches that limit. Can be increased later;
  // dispenso does not guarantee ABI stability.
  static constexpr int32_t kMaxCpus = 1024;
  static constexpr int32_t kBitsPerWord = 64;
  static constexpr int32_t kNumWords = kMaxCpus / kBitsPerWord;
  uint64_t words_[kNumWords];
#endif
};
 
struct ThreadGroup {
  std::vector<int32_t> cpus; 
  CpuSet affinityMask; 
};
 
namespace detail {
DISPENSO_DLL_ACCESS CpuSet parseLinuxCpuList(const char* input);
} // namespace detail
 
} // namespace dispenso