fix(partition): 🐛 place partitions inside the cgroup the batch system carved

cpp/monoprop/detail/partition/CpuTopology.cpp

@@ -15,7 +15,10 @@
 #include "monoprop/detail/partition/CpuTopology.h"
 
 #include <algorithm>
+#include <cstdio>
 #include <map>
+#include <mutex>
+#include <print>
 #include <vector>
 
 #include <hwloc.h>
@@ -88,7 +91,7 @@
 
 auto placement_order(const std::vector<PhysicalCore> &cores, size_t n, size_t group_index, size_t group_count)
     -> std::vector<int> {
-    if (cores.empty() || group_count * n > cores.size()) {
+    if (cores.empty() || group_count == 0 || group_count * n > cores.size()) {
         return {};
     }
 
@@ -141,7 +144,7 @@
     if (offset + n > order.size()) {
         return {};
     }
     return std::vector<int>(order.begin() + static_cast<std::ptrdiff_t>(offset),
                             order.begin() + static_cast<std::ptrdiff_t>(offset + n));
 }
 
@@ -150,7 +153,7 @@
 /* ── enumerate_physical_cores ──────────────────────────────────────────────── */
 
 auto enumerate_physical_cores() -> std::vector<PhysicalCore> {
     const auto topo = get_topology();
     if (!topo) {
         return {};
     }
@@ -172,7 +175,7 @@
 
     const unsigned num_cores = hwloc_get_nbobjs_by_depth(topo, core_depth);
     for (unsigned i = 0; i < num_cores; ++i) {
         const hwloc_obj_t core = hwloc_get_obj_by_depth(topo, core_depth, i);
         if (!core || !core->cpuset) {
             continue;
         }
@@ -199,7 +202,7 @@
          * receive their own singleton domain so the placement algorithm can still spread across
          * whatever structure the topology does have. */
         int domain;
         const hwloc_obj_t l3 = hwloc_get_ancestor_obj_by_type(topo, HWLOC_OBJ_L3CACHE, core);
         if (l3) {
             const auto [it, inserted] = l3_domain_map.emplace(l3->logical_index, next_domain_id);
             if (inserted) {
@@ -211,22 +214,96 @@
             domain = next_domain_id++;
         }
 
         cores.push_back(PhysicalCore{rep, domain});
     }
 
     hwloc_bitmap_free(allowed);
     return cores;
 }
 
+/* ── affinity_mask_words ───────────────────────────────────────────────────── */
+
+auto affinity_mask_words(uint64_t *out, size_t nwords) -> bool {
+    if (out == nullptr || nwords == 0) {
+        return false;
+    }
+    std::fill_n(out, nwords, uint64_t{0});
+    const auto topo = get_topology();
+    if (!topo) {
+        return false;
+    }
+    const hwloc_cpuset_t allowed = effective_allowed_cpuset(topo);
+    if (!allowed) {
+        return false;
+    }
+    // Refused rather than truncated: a truncated mask could compare disjoint against a peer it overlaps.
+    const int last = hwloc_bitmap_last(allowed);
+    const bool representable = last >= 0 && static_cast<size_t>(last) < nwords * 64;
+    if (representable) {
+        for (int pu = hwloc_bitmap_first(allowed); pu >= 0; pu = hwloc_bitmap_next(allowed, pu)) {
+            out[static_cast<size_t>(pu) / 64] |= uint64_t{1} << (static_cast<size_t>(pu) % 64);
+        }
+    }
+    hwloc_bitmap_free(allowed);
+    return representable;
+}
+
+/* ── masks_are_pairwise_disjoint ───────────────────────────────────────────── */
+
+auto masks_are_pairwise_disjoint(const uint64_t *masks, size_t n, size_t words) -> bool {
+    if (masks == nullptr || words == 0 || n < 2) {
+        return false;
+    }
+    // An all-zero mask is disjoint from everything, so empty is rejected before the pairwise test.
+    for (size_t r = 0; r < n; ++r) {
+        bool any = false;
+        for (size_t w = 0; w < words && !any; ++w) {
+            any = masks[(r * words) + w] != 0;
+        }
+        if (!any) {
+            return false;
+        }
+    }
+    for (size_t a = 0; a < n; ++a) {
+        for (size_t b = a + 1; b < n; ++b) {
+            for (size_t w = 0; w < words; ++w) {
+                if ((masks[(a * words) + w] & masks[(b * words) + w]) != 0) {
+                    return false; // two peers share a CPU: not private
+                }
+            }
+        }
+    }
+    return true;
+}
+
 /* ── partition_cpusets ─────────────────────────────────────────────────────── */
 
-auto partition_cpusets(size_t n, size_t group_index, size_t group_count) -> std::vector<CpuSet> {
+auto partition_cpusets(size_t n, size_t group_index, size_t group_count, bool mask_is_private) -> std::vector<CpuSet> {
     if (!config::get().partition_pinning) {
         return {};
     }
     const auto cores = enumerate_physical_cores();
+
+    // A private mask IS this rank's share: the launcher already separated co-located ranks.
+    if (mask_is_private) {
+        group_index = 0;
+        group_count = 1;
+    }
     const auto order = topo_detail::placement_order(cores, n, group_index, group_count);
 
+    if (order.empty()) {
+        static std::once_flag warned;
+        std::call_once(warned, [&] {
+            std::print(stderr,
+                       "monoprop: partition pinning requested but not possible "
+                       "({} cores visible, {} groups x {} partitions); threads run unpinned.\n",
+                       cores.size(),
+                       group_count,
+                       n);
+            std::fflush(stderr);
+        });
+    }
+
     std::vector<CpuSet> sets(order.size());
     for (size_t i = 0; i < order.size(); ++i) {
         sets[i] = CpuSet{order[i]};
@@ -240,7 +317,7 @@
     if (set.pu < 0) {
         return;
     }
     const auto topo = get_topology();
     if (!topo) {
         return;
     }

cpp/monoprop/detail/partition/CpuTopology.h

@@ -24,7 +24,9 @@
 
 #pragma once
 
+#include <climits>
 #include <cstddef>
+#include <cstdint>
 #include <vector>
 
 #include "monoprop/detail/EnvConfig.h"
@@ -88,6 +90,20 @@ auto placement_order(const std::vector<PhysicalCore> &cores, size_t n, size_t gr
  */
 auto enumerate_physical_cores() -> std::vector<PhysicalCore>;
 
+//! Affinity-mask exchange width, in 64-bit words. A mask needing more is "cannot classify", never private.
+inline constexpr size_t kAffinityMaskWords = 64;
+
+static_assert(kAffinityMaskWords > 0 && kAffinityMaskWords <= static_cast<size_t>(INT_MAX),
+              "the affinity-mask width is an MPI_Allgather element count, which is an int");
+
+//! This process's allowed cpuset as @p nwords 64-bit words; false with @p out zeroed when it does not fit.
+auto affinity_mask_words(uint64_t *out, size_t nwords) -> bool;
+
+/*! @brief Whether the @p n masks of @p words words laid end to end in @p masks are pairwise disjoint.
+ *  False for @p n < 2 and for any empty mask: all-zero is disjoint from everything, and shared is the safe error.
+ */
+[[nodiscard]] auto masks_are_pairwise_disjoint(const uint64_t *masks, size_t n, size_t words) -> bool;
+
 /*!
  * @brief Build placement tokens for one MPI rank's partitions.
  *
@@ -98,10 +114,15 @@ auto enumerate_physical_cores() -> std::vector<PhysicalCore>;
  * @param n            Number of partitions to place.
  * @param group_index  This rank's 0-based index among the co-located ranks on the host.
  * @param group_count  Total number of co-located ranks on the host.
+ * @param mask_is_private  True only when the co-located ranks' affinity masks have been measured
+ *                     pairwise DISJOINT (PartitionGroup::classify_node_masks_), so this mask is our share.
  * @returns Vector of @p n CpuSet tokens, or empty when @c monoprop_PARTITION_PINNING is disabled,
- *          hwloc is unavailable, or the host cannot provide @p group_count × @p n distinct cores.
+ *          hwloc is unavailable, or fewer than @p group_count x @p n cores are visible (@p n when private).
+ *
+ * @note Under @p mask_is_private the group split is skipped: our share is already this rank's alone.
  */
-auto partition_cpusets(size_t n, size_t group_index = 0, size_t group_count = 1) -> std::vector<CpuSet>;
+auto partition_cpusets(size_t n, size_t group_index = 0, size_t group_count = 1, bool mask_is_private = false)
+    -> std::vector<CpuSet>;
 
 /*!
  * @brief Bind the calling thread to the PU identified by @p set.

cpp/monoprop/detail/partition/PartitionGroup.h

@@ -14,8 +14,10 @@
 
 #pragma once
 
+#include <array>
 #include <condition_variable>
 #include <cstddef>
+#include <cstdint>
 #include <exception>
 #include <functional>
 #include <memory>
@@ -59,7 +61,7 @@ class PartitionGroup {
           errs_(static_cast<size_t>(n_partitions)) {
         make_transport_();
         discover_node_peers_();
-        cpusets_ = topo_partition_cpusets(n_, node_rank_, node_size_);
+        cpusets_ = topo_partition_cpusets(n_, node_rank_, node_size_, node_mask_private_);
         start_masters_();
         // The masters are already running, so a ctor throw must not escape: ~PartitionGroup would never run,
         // and destroying joinable threads during unwinding calls std::terminate.
@@ -79,10 +81,11 @@ class PartitionGroup {
           parent_(src.parent_),
           node_rank_(src.node_rank_),
           node_size_(src.node_size_),
+          node_mask_private_(src.node_mask_private_),
           partitions_(static_cast<size_t>(src.n_)),
           errs_(static_cast<size_t>(src.n_)) {
         make_transport_();
-        cpusets_ = topo_partition_cpusets(n_, node_rank_, node_size_);
+        cpusets_ = topo_partition_cpusets(n_, node_rank_, node_size_, node_mask_private_);
         start_masters_();
         try { // see the primary ctor: a throw past live masters would std::terminate
             run_on_all([&](int r) {
@@ -146,11 +149,12 @@ class PartitionGroup {
     }
 
     // Free-function wrapper so the header compiles on non-Linux (where partition_cpusets returns {}).
-    static auto topo_partition_cpusets(int n, int group_index, int group_count)
+    static auto topo_partition_cpusets(int n, int group_index, int group_count, bool mask_is_private)
         -> std::vector<monoprop::detail::partition::CpuSet> {
         return monoprop::detail::partition::partition_cpusets(static_cast<size_t>(n),
                                                               static_cast<size_t>(group_index),
-                                                              static_cast<size_t>(group_count));
+                                                              static_cast<size_t>(group_count),
+                                                              mask_is_private);
     }
 
     // Under an MPI parent, find how many ranks share this host and which we are, so each co-located rank
@@ -162,11 +166,31 @@ class PartitionGroup {
             MPI_Comm_split_type(parent_.mpi, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &node);
             MPI_Comm_rank(node, &node_rank_);
             MPI_Comm_size(node, &node_size_);
+            classify_node_masks_(node);
             MPI_Comm_free(&node);
         }
 #endif
     }
 
+#ifdef monoprop_ENABLE_MPI
+    // A rank seeing 16 of 128 CPUs is equally "my own 16" and "eight of us share these 16": only the masks tell.
+    auto classify_node_masks_(MPI_Comm node) -> void {
+        node_mask_private_ = false;
+        if (node_size_ <= 1) {
+            return; // nobody to collide with; the normal split already handles group_count == 1
+        }
+        constexpr size_t kMaskWords = monoprop::detail::partition::kAffinityMaskWords;
+        std::array<uint64_t, kMaskWords> mine{};
+        // Refusal zeroes `mine` and an all-zero row is never private, so no reduction of the verdict is needed.
+        monoprop::detail::partition::affinity_mask_words(mine.data(), kMaskWords);
+        std::vector<uint64_t> all(kMaskWords * static_cast<size_t>(node_size_), 0);
+        MPI_Allgather(mine.data(), kMaskWords, MPI_UINT64_T, all.data(), kMaskWords, MPI_UINT64_T, node);
+        node_mask_private_ = monoprop::detail::partition::masks_are_pairwise_disjoint(all.data(),
+                                                                                      static_cast<size_t>(node_size_),
+                                                                                      kMaskWords);
+    }
+#endif
+
     auto make_transport_() -> void {
 #ifdef monoprop_ENABLE_MPI
         if (parent_.kind == mpi::Comm::Kind::Mpi && mpi::size(parent_) > 1) {
@@ -253,6 +277,7 @@ class PartitionGroup {
     mpi::Comm parent_;                  // enclosing communicator (size R) — decides the transport
     int node_rank_ = 0;                 // this rank's index among the ranks sharing the host
     int node_size_ = 1;                 // how many parent ranks share the host (1 unless MPI R>1)
+    bool node_mask_private_ = false;    // set by classify_node_masks_; copied, never re-derived, by the copy ctor
     std::unique_ptr<mpi::ShmComm> shm_; // set iff R == 1
 #ifdef monoprop_ENABLE_MPI
     std::unique_ptr<mpi::HybridComm> hyb_; // set iff R > 1