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 <string>
#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>
#if defined(__FreeBSD__)
#include <pthread_np.h>
#include <sys/cpuset.h>
#include <sys/param.h>
#endif
#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 uses pthread_setaffinity_np, as on Linux. NUMA and cache
// topology are read from sysctl (vm.ndomains, kern.sched.topology_spec); when
// the kernel reports no cache levels, cache group queries return empty and
// buildThreadGroups() falls back to contiguous chunking.
 
#ifndef DISPENSO_MAX_GROUP_SIZE
#define DISPENSO_MAX_GROUP_SIZE 16
#endif
 
namespace dispenso {
 
namespace detail {
#if defined(DISPENSO_CPUSET_LINUXY)
// The native affinity-set type: cpu_set_t on Linux, cpuset_t on FreeBSD. Aliased
// in detail rather than as a global `typedef cpuset_t cpu_set_t` so that
// including this header does not inject a cpu_set_t name into every translation
// unit (which could clash with a consumer's own definition).
#if defined(__FreeBSD__)
using CpuSetNative = ::cpuset_t;
#else
using CpuSetNative = ::cpu_set_t;
#endif
#endif // DISPENSO_CPUSET_LINUXY
} // namespace detail
 
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)
  detail::CpuSetNative 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);
 
DISPENSO_DLL_ACCESS std::vector<CacheGroup> parseCacheGroupsFromTopologySpec(
    const std::string& xml,
    int cacheIndex);
 
DISPENSO_DLL_ACCESS std::vector<ThreadGroup> buildGroupsFromCacheTopology(
    const std::vector<CacheGroup>& l2Groups,
    const std::vector<CacheGroup>& l3Groups,
    int32_t maxGroupSize);
} // namespace detail
 
} // namespace dispenso