Rewrite nvtaskset and implementation of partitioning for unmodified tasks

Rather than requiring libsmctrl.so to be preloaded, we now wrap
libcuda.so.1. All CUDA-using applications will load libcuda.so.1,
ensuring that our wrapper will always be dynamically loaded, no
matter if LD_PRELOAD is enabled, or if a program has been staticly
linked. All that needs to be done is that the location of our
"fake" libcuda.so.1 need to be put within the loader search path.
This can be done by setting LD_LIBRARY_PATH, or by installing
our wrapper into /lib/x86_64-linux-gnu.

The mask can still be set via the LIBSMCTRL_MASK environment
variable, but the easier-to-use nvtaskset tool is now the
recommended way to view or change the supreme TPC mask for any
CUDA-using application. This allows launching a program on the
first two GPCs via a command as simple as:
    ./nvtaskset -g 0-1 ./a_program a_program_args
(Note that use of the -g option requires the nvdebug kernel module
to first be loaded.)

These changes support the final version of the ECRTS'25 paper.

Note that nvtaskset does not yet fully support multi-GPU systems.

Bugfixes:
- Fix crash that would occur if both libsmctrl.so and libsmctrl.a
  were built into an application.
- Correctly use GPU ID when initializing a context in
  `libsmctrl_test_gpc_info`.
- Include `nvtaskset` as a prerequisite for
  `libsmctrl_test_supreme_mask`.
- Fix malfunction of `libsmctrl_test_gpc_info` if
  CUDA_VISIBLE_DEVICES is set.

Other minor changes:
- Adds make target to run all the tests.
- Fixes typos in comments.
- Enables -Wall build option.
- Upgrades supreme mask from 64 to 128 bits.
- Removes `detect_parker_soc()` from the global namespace.
- Adjusts test messages to be more succinct.
- Updates README with overview of how to partition unmodified
  applications, more details on the tests, and information on the
  new ECRTS'25 paper.

1 files changed, 347 insertions, 173 deletions

diff --git a/nvtaskset.c b/nvtaskset.c
index 74f88d7..4901cbe 100644
--- a/nvtaskset.c
+++ b/nvtaskset.c
@@ -1,210 +1,384 @@
 // Copyright 2025 Joshua Bakita
-// taskset-like utility for the GPU
+// Show or change the GPU core affinity for a CUDA process
+// taskset-like utility for NVIDIA GPUs
 #define _GNU_SOURCE // For program_invocation_name
 #include <argp.h>
+#include <dirent.h>
 #include <errno.h>
 #include <error.h>
+#include <fcntl.h>
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <sys/ipc.h>
-#include <sys/shm.h>
-#include <sys/types.h>
+#include <sys/mman.h>
 #include <unistd.h>
 
 #include <cuda.h> // To help with getting GPC info
 
 #include "libsmctrl.h"
 
-const char* maintainer = "<jbakita@cs.unc.edu>";
-const char* version = "nvtaskset 2025.03";
-const char* desc = "taskset-like utility for NVIDIA GPUs.";
+#define LINK_NAME "/memfd:libsmctrl"
+
+// TODO: Write automated tests:
+// - Change region of non-existent PID
+// - Change region of permission denied PID
+// - Change region of non-GPU PID
+// - Change GPC list
+// - Change TPC list
+// - Change TPC mask
+// - Start TPC mask
+// - Start TPC list
+// - Start GPC list
+// - Start with subargument containing -
+// - Query GPC list
+// - Query TPC list
+// - Query TPC mask
+// - Set GPC list w/ non-existant GPC
+// - Set TPC list w/ non-existant TPC
+
+// Private symbols from libsmctrl
+extern bool libsmctrl_is_mps_running();
+extern uint128_t strtou128(const char *nptr, char **endptr, int base);
+
+const char *argp_program_bug_address = "<jbakita@cs.unc.edu>";
+const char *argp_program_version = "nvtaskset 2025.06";
+const char desc[] = "Show or change the GPU core affinity for a CUDA process\v"
+                    "Warning: When using GPC lists, this tool currently "
+                    "derives TPC to GPC mappings from the first NVIDIA GPU in "
+                    "the system (by PCI bus ID) device. To use the mappings "
+                    "for a different device, use the `libsmctrl_test_get_info` "
+                    "tool to get the bitmask of TPCs associated with each GPC, "
+                    "OR them, and then set that bitmask via this tool. Better "
+                    "multi-GPU support is intended for a future release.\n\n"
+                    "Inspired by the Linux taskset utility.";
+const char args_doc[] = "[mask | list] [pid | cmd [args...]]";
 
 const struct argp_option opts[] = {
+        {"gpc-list", 'g', NULL, 0, "Specify partition as a list of GPCs"},
+        {"tpc-list", 't', NULL, 0, "Specify partition as a list of TPCs"},
+        {"pid",      'p', NULL, 0, "Operate on an existing PID"},
         {0}
 };
 
-unsigned __int128 strtou128(const char *nptr, char **endptr, int base) {
-        unsigned __int128 result = 0;
-        if (base != 16)
-                error(1, EINVAL, "Internal error");
-        // Skip a "0x" prefix. Safe due to early evaluation
-        if (*nptr == '0' && (*(nptr + 1) == 'x' || *(nptr + 1) == 'X'))
-                nptr += 2;
-        // Until hitting an invalid character
-        while (1) {
-                if (*nptr >= 'a' && *nptr <= 'f')
-                        result = result << 4 | (*nptr - 'a' + 10);
-                else if (*nptr >= 'A' && *nptr <= 'F')
-                        result = result << 4 | (*nptr - 'A' + 10);
-                else if (*nptr >= '0' && *nptr <= '9')
-                        result = result << 4 | (*nptr - '0');
-                else
-                        break;
-                nptr++;
-        }
-        if (endptr)
-                *endptr = (char*)nptr;
-        return result;
-}
-
+// Create a CUDA context and query the associated GPC to TPC mappings
+// Based off logic in libsmctrl_test_gpc_info
 void libsmctrl_get_gpc_info_ext_easy(uint32_t* num_gpcs, uint128_t** masks, int gpu_id) {
         int res;
         CUcontext ctx;
-        // XXX: Copied from libsmctrl_test_gpc_info
-        // Tell CUDA to use PCI device id ordering (to match nvdebug)
-        putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID");
-        // A CUDA context is required before reading the topology information
-        if ((res = cuInit(0))) {
-                const char* name;
-                cuGetErrorName(res, &name);
-                fprintf(stderr, "%s: Unable to initialize CUDA, error %s\n", program_invocation_name, name);
-                exit(1);
-        }
-        if ((res = cuCtxCreate(&ctx, 0, 0))) {
-                const char* name;
-                cuGetErrorName(res, &name);
-                fprintf(stderr, "%s: Unable to create a CUDA context, error %s\n", program_invocation_name, name);
-                exit(1);
-        }
-        // Pull topology information from libsmctrl
-        if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) {
-                error(0, res, "libsmctrl_get_gpc_info() failed");
-                if (res == ENOENT)
-                        fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name);
-                if (res == EIO)
-                        fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name);
-                exit(1);
-        }
-        // Not copied
+        char *old_order = NULL;
+        // Tell CUDA to use PCI device id ordering (to match nvdebug)
+        putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID");
+        // Allow CUDA to see all devices (to better match nvdebug)
+        if (getenv("CUDA_VISIBLE_DEVICES")) {
+                if (!(old_order = strdup(getenv("CUDA_VISIBLE_DEVICES"))))
+                        error(1, errno, "Unable to allocate environment string");
+                unsetenv("CUDA_VISIBLE_DEVICES");
+        }
+        // A CUDA context is required before reading the topology information
+        if ((res = cuInit(0))) {
+                const char* name;
+                cuGetErrorName(res, &name);
+                error(1, 0, "Unable to create a initialize CUDA, error %s\n", name);
+        }
+        if ((res = cuCtxCreate(&ctx, 0, gpu_id))) {
+                const char* name;
+                cuGetErrorName(res, &name);
+                error(1, 0, "Unable to create a CUDA context, error %s\n", name);
+        }
+        // Pull topology information from libsmctrl
+        if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) {
+                error(0, res, "libsmctrl_get_gpc_info() failed");
+                if (res == ENOENT)
+                        fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name);
+                if (res == EIO)
+                        fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name);
+                exit(1);
+        }
+        // Restore the environment (in case we exec() later)
         unsetenv("CUDA_DEVICE_ORDER");
+        if (old_order) {
+                setenv("CUDA_VISIBLE_DEVICES", old_order, 1);
+                free(old_order);
+        }
 }
 
-int main(int argc, char **argv) {
-        if (argc < 3) {
-                fprintf(stderr, "Usage: %s -p <hex mask> <pid>\n", argv[0]);
-                fprintf(stderr, "       %s <hex mask> <command> <argument...>\n", argv[0]);
-                fprintf(stderr, "       %s --gpc-list <gpc list> <command> <argument...>\n", argv[0]);
-                fprintf(stderr, " <hex mask> has a bit set for each TPC to be enabled\n");
-                return 1;
-        }
-        // TODO: Use a proper argument parser
-        if (strcmp("-p", argv[1]) == 0) { // Setting mask on running task
-                char *end;
-                pid_t target_pid = strtoul(argv[2], &end, 10);
-                // strtoul stores a pointer to the first invalid character in `end`
-                if (*end != '\0') {
-                        fprintf(stderr, "Invalid character \"%c\" in PID argument.\n", *end);
-                        return 1;
+int parse_list(bool use_gpcs, char* list, uint128_t *mask_out) {
+        // We support the same ranges as taskset, e.g., X,Y,Z and X,Y-Z
+        uint32_t num_xpcs = 0; // Either TPC or GPC count, i.e., "X"PC
+        uint128_t* masks = NULL;
+        // TODO: Allow specifying GPU ID, rather than assuming 0!
+        if (use_gpcs)
+                libsmctrl_get_gpc_info_ext_easy(&num_xpcs, &masks, 0);
+        else
+                libsmctrl_get_tpc_info_cuda(&num_xpcs, 0);
+        uint128_t mask = 0;
+        int range_start_xpc = -1;
+        char* start = list;
+        int len = strlen(list);
+        // Convert comma-seperated GPC/TPC list into a mask
+        for (int i = 0; i < len + 1; i++) {
+                if (list[i] == ',' || list[i] == '\0') {
+                        list[i] = '\0';
+                        int xpc = atoi(start);
+                        if (xpc > num_xpcs - 1)
+                                error(1, EINVAL, "%s is not a valid %s ID", start, use_gpcs ? "GPC" : "TPC");
+                        // Handle ranges
+                        if (range_start_xpc != -1) {
+                                if (range_start_xpc >= xpc)
+                                        error(1, EINVAL, "Malformed %s range", use_gpcs ? "GPC" : "TPC");
+                                while (range_start_xpc <= xpc) {
+                                        if (use_gpcs)
+                                                mask |= masks[range_start_xpc];
+                                        else
+                                                mask |= (uint128_t)1 << range_start_xpc;
+                                        range_start_xpc++;
+                                }
+                                range_start_xpc = -1;
+                        } else {
+                                if (use_gpcs)
+                                        mask |= masks[xpc];
+                                else
+                                        mask |= (uint128_t)1 << xpc;
+                        }
+                        start = list + i + 1;