Rework user-defined mapping. (#114)

Signed-off-by: Samuel K. Gutierrez <[email protected]>

src/qvi-hwloc.cc

@@ -1388,26 +1388,18 @@ static int
 split_cpuset_chunk_size(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
-    int npieces,
-    int *chunk
+    uint_t npieces,
+    uint_t *chunk_size
 ) {
-    int rc = QV_SUCCESS;
-
-    if (npieces <= 0) {
-        rc = QV_ERR_INVLD_ARG;
-        goto out;
-    }
-
-    int npus;
-    rc = qvi_hwloc_get_nobjs_in_cpuset(
+    int npus = 0;
+    int rc = qvi_hwloc_get_nobjs_in_cpuset(
         hwl, QV_HW_OBJ_PU, cpuset, &npus
     );
-out:
-    if (rc != QV_SUCCESS) {
-        *chunk = 0;
+    if (rc != QV_SUCCESS || npieces == 0 || npus == 0 ) {
+        *chunk_size = 0;
     }
     else {
-        *chunk = npus / npieces;
+        *chunk_size = npus / npieces;
     }
     return rc;
 }
@@ -1419,8 +1411,8 @@ static int
 split_cpuset_by_range(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
-    int base,
-    int extent,
+    uint_t base,
+    uint_t extent,
     hwloc_bitmap_t result
 ) {
     // Zero-out the result bitmap that will encode the split.
@@ -1434,7 +1426,7 @@ split_cpuset_by_range(
     );
     if (rc != QV_SUCCESS) return rc;
     // Calculate split based on given range.
-    for (int i = base; i < base + extent; ++i) {
+    for (uint_t i = base; i < base + extent; ++i) {
         hwloc_obj_t dobj;
         rc = qvi_hwloc_get_obj_in_cpuset_by_depth(
             hwl, cpuset, pu_depth, i, &dobj
@@ -1450,13 +1442,12 @@ split_cpuset_by_range(
     return rc;
 }
 
-// TODO(skg) Merge with split_cpuset_by_range().
 int
 qvi_hwloc_get_cpuset_for_nobjs(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
     qv_hw_obj_type_t obj_type,
-    int nobjs,
+    uint_t nobjs,
     hwloc_cpuset_t *result
 ) {
     hwloc_bitmap_t iresult = nullptr;
@@ -1469,7 +1460,7 @@ qvi_hwloc_get_cpuset_for_nobjs(
     );
     if (rc != QV_SUCCESS) goto out;
     // Calculate cpuset based on number of desired objects.
-    for (int i = 0; i < nobjs; ++i) {
+    for (uint_t i = 0; i < nobjs; ++i) {
         hwloc_obj_t dobj;
         rc = qvi_hwloc_get_obj_in_cpuset_by_depth(
             hwl, cpuset, obj_depth, i, &dobj
@@ -1491,33 +1482,26 @@ qvi_hwloc_get_cpuset_for_nobjs(
 }
 
 int
-qvi_hwloc_split_cpuset_by_color(
+qvi_hwloc_split_cpuset_by_chunk_id(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
-    int ncolors,
-    int color,
+    uint_t nchunks,
+    uint_t chunk_id,
     hwloc_cpuset_t result
 ) {
-    int chunk = 0;
+    uint_t chunk_size = 0;
     int rc = split_cpuset_chunk_size(
-        hwl, cpuset, ncolors, &chunk
+        hwl, cpuset, nchunks, &chunk_size
     );
     if (rc != QV_SUCCESS) return rc;
-    // This happens when n > npus. We can't support that split.
-    // TODO(skg) Perhaps we can create an empty cpuset that denotes no
-    // resources?
-    if (chunk == 0) {
-        return QV_ERR_SPLIT;
-    }
-    // Group IDs must be < n: 0, 1, ... , ncolors-1.
-    // TODO(skg) We could also allow this. That might mean that this processor
-    // should not be considered in the split.
-    if (color >= ncolors) {
+    // 0 chunk_size likely caused by nonsensical split request.
+    // Chunk IDs must be < nchunks: 0, 1, ... , nchunks-1.
+    if (chunk_size == 0 || chunk_id >= nchunks) {
         return QV_ERR_SPLIT;
     }
 
     return split_cpuset_by_range(
-        hwl, cpuset, chunk * color, chunk, result
+        hwl, cpuset, chunk_size * chunk_id, chunk_size, result
     );
 }
 

src/qvi-hwloc.h

@@ -376,11 +376,11 @@ qvi_hwloc_get_device_in_cpuset(
  *
  */
 int
-qvi_hwloc_split_cpuset_by_color(
+qvi_hwloc_split_cpuset_by_chunk_id(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
-    int ncolors,
-    int color,
+    uint_t nchunks,
+    uint_t chunk_id,
     hwloc_cpuset_t result
 );
 
@@ -392,7 +392,7 @@ qvi_hwloc_get_cpuset_for_nobjs(
     qvi_hwloc_t *hwl,
     hwloc_const_cpuset_t cpuset,
     qv_hw_obj_type_t obj_type,
-    int nobjs,
+    uint_t nobjs,
     hwloc_cpuset_t *result
 );
 
@@ -427,7 +427,7 @@ struct qvi_hwloc_bitmap_s {
     {
         qvim_rc = qvi_hwloc_bitmap_calloc(&data);
         if (qvim_rc != QV_SUCCESS) return;
-        set(bitmap);
+        qvim_rc = set(bitmap);
     }
     /** Copy constructor. */
     qvi_hwloc_bitmap_s(const qvi_hwloc_bitmap_s &src)

src/qvi-map.cc

@@ -125,6 +125,41 @@ qvi_map_fid_mapped(
     return map.find(cid) != map.end();
 }
 
+int
+qvi_map_colors(
+    qvi_map_t &map,
+    const std::vector<int> &fcolors,
+    const qvi_hwloc_cpusets_t &tres
+) {
+    // Note: the array index i of fcolors is the color requested by task i.
+    // Determine the number of distinct colors provided in the colors array.
+    std::set<int> color_set(fcolors.begin(), fcolors.end());
+    const uint_t nfrom = color_set.size();
+    // For convenience, we convert the set to a vector for later use.
+    std::vector<int> color_vec(color_set.begin(), color_set.end());
+    // Maps a given color to its corresponding vector index. For example:
+    // colors = {3, 5, 3, 4}, color_set = {3, 4, 5}, color_vec = {3, 4, 5}
+    // color_set_index (csi) = {0, 1, 2}, since we have three distinct colors.
+    // color2csi = {3:0, 4:1, 5:2}
+    qvi_map_t color2csi;
+    for (uint_t i = 0; i < color_vec.size(); ++i) {
+        color2csi.insert({color_vec[i], i});
+    }
+    // Create a mapping of color_set indices to cpuset indices.
+    qvi_map_t csi2cpui;
+    int rc = qvi_map_packed(csi2cpui, nfrom, tres);
+    if (rc != QV_SUCCESS) return rc;
+    // Now map the task colors to their respective cpusets.
+    for (uint_t fid = 0; fid < fcolors.size(); ++fid) {
+        // Already mapped (potentially by some other mapper).
+        if (qvi_map_fid_mapped(map, fid)) continue;
+        const int csi = color2csi.at(fcolors[fid]);
+        const int tid = csi2cpui.at(csi);
+        map.insert({fid, tid});
+    }
+    return rc;
+}
+
 int
 qvi_map_packed(
     qvi_map_t &map,
@@ -145,7 +180,7 @@ qvi_map_packed(
             // Already mapped (potentially by some other mapper).
             if (qvi_map_fid_mapped(map, fid)) continue;
             // Else map the consumer to the resource ID.
-            map.insert(std::make_pair(fid, tid));
+            map.insert({fid, tid});
             nmapped++;
         }
     }
@@ -169,7 +204,7 @@ qvi_map_spread(
         // Already mapped (potentially by some other mapper).
         if (qvi_map_fid_mapped(map, fid)) continue;
         // Mod to loop around 'to resource' IDs.
-        map.insert(std::make_pair(fid, (tid++) % ntres));
+        map.insert({fid, (tid++) % ntres});
         nmapped++;
     }
     return QV_SUCCESS;
@@ -188,7 +223,7 @@ qvi_map_disjoint_affinity(
             // Already mapped (potentially by some other mapper).
             if (qvi_map_fid_mapped(map, fid)) continue;
             // Map the consumer ID to its resource ID.
-            map.insert(std::make_pair(fid, tid));
+            map.insert({fid, tid});
         }
     }
     return QV_SUCCESS;

src/qvi-map.h

@@ -112,6 +112,16 @@ qvi_map_disjoint_affinity(
     const qvi_map_shaffinity_t &damap
 );
 
+/**
+ * Performs a mapping between the provided colors to the provided cpusets.
+ */
+int
+qvi_map_colors(
+    qvi_map_t &map,
+    const std::vector<int> &fcolors,
+    const qvi_hwloc_cpusets_t &tres
+);
+
 /**
  * Performs an affinity preserving mapping.
  */