1 files changed, 210 insertions, 0 deletions

diff --git a/nvtaskset.c b/nvtaskset.c
new file mode 100644
index 0000000..b8873ab
--- /dev/null
+++ b/nvtaskset.c
@@ -0,0 +1,210 @@
+// Copyright 2025 Joshua Bakita
+// taskset-like utility for the GPU
+#define _GNU_SOURCE // For program_invocation_name
+#include <argp.h>
+#include <errno.h>
+#include <error.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 <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.";
+
+const struct argp_option opts[] = {
+        {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;
+}
+
+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
+        unsetenv("CUDA_DEVICE_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;
+                }
+                unsigned __int128 mask = strtou128(argv[3], &end, 16);
+                if (*end != '\0') {
+                        fprintf(stderr, "Invalid character \"%c\" in mask argument.\n", *end);
+                        return 1;
+                }
+                // The shared memory lookup key is the lower 16-bits of the PID | "sm"
+                key_t shm_key = target_pid << 16 | (int)'s' << 8 | (int) 'm';
+                // Get a handle to the 128-bit shared memory region
+                int shmid = shmget(shm_key, 16, 0);
+                if (shmid == -1)
+                        error(1, errno, "Unable to find control region for PID %d", target_pid);
+                // Open the shared memory region
+                unsigned __int128 *supreme_mask = shmat(shmid, NULL, 0);
+                if (supreme_mask == (void*)-1)
+                        error(1, errno, "Unable to open control region for PID %d", target_pid);
+                // Write the requested mask into the shared memory region
+                *supreme_mask = mask;
+        } else { // Starting a new task with a mask
+                // TODO: Check other locations for nvidia-cuda-mps-control if its not on the path
+                // TODO: Use dup2() to redirect MPS startup messages
+                int ret = system("echo -n | nvidia-cuda-mps-control");
+                if (ret == -1)
+                        error(1, errno, "Unable to run subshell to check MPS status");
+                if (ret != 0) { // Control deamon not yet started
+                        fprintf(stderr, "nvtaskset: MPS control deamon does not appear to be running. Automatically starting...\n");
+                        ret = system("nvidia-cuda-mps-control -d");
+                        if (ret == -1)
+                                error(1, errno, "Unable to run subshell to start MPS");
+                        if (ret == 1) {
+                                fprintf(stderr, "nvtaskset: Error starting MPS control deamon. Terminating...\n");
+                                return 1;
+                        }
+                        fprintf(stderr, "nvtaskset: Done. Use \"echo quit | nvidia-cuda-mps-control\" to terminate it later as desired.\n");
+                }
+                // Tell loader to initialize libsmctrl.so first
+                // TODO: Append, rather than overwrite LD_PRELOAD
+                setenv("LD_PRELOAD", "./libsmctrl.so", 1);
+                // Explictly set the number of channels, otherwise CUDA will only use two
+                // (see paper for why that causes problems)
+                setenv("CUDA_DEVICE_MAX_CONNECTIONS", "8", 1);
+                // Check if a mask, or a list of GPCs is being provided
+                if (strcmp(argv[1], "--gpc-list") == 0) {
+                        // TODO: Support the full syntax that taskset supports
+                        // We just support X,Y,Z for now
+                        uint32_t num_gpcs = 0;
+                        uint128_t* masks = NULL;
+                        // TODO: Allow specifying GPU ID, rather than assuming 0!
+                        libsmctrl_get_gpc_info_ext_easy(&num_gpcs, &masks, 0);
+                        uint128_t mask = 0;
+                        int range_start_gpc = -1;
+                        char* start = argv[2];
+                        int len = strlen(argv[2]);
+                        // TODO: Handle invalid input cleanly.
+                        // Convert comma-seperated GPC list into a mask
+                        for (int i = 0; i < len + 1; i++) {
+                                if (argv[2][i] == ',' || argv[2][i] == '\0') {
+                                        argv[2][i] = '\0';
+                                        int gpc = atoi(start);
+                                        if (gpc > num_gpcs - 1) {
+                                                fprintf(stderr, "Invalid GPC ID '%s'!\n", start);
+                                                return 1;
+                                        }
+                                        // Handle ranges
+                                        if (range_start_gpc != -1) {
+                                                if (range_start_gpc >= gpc) {
+                                                        fprintf(stderr, "Invalid GPC range!\n");
+                                                        return 1;
+                                                }
+                                                while (range_start_gpc <= gpc) {
+                                                        //printf("gpc %i\n", range_start_gpc);
+                                                        mask |= masks[range_start_gpc];
+                                                        range_start_gpc++;
+                                                }
+                                                range_start_gpc = -1;
+                                        } else {
+                                                //printf("gpc %i\n", gpc);
+                                                mask |= masks[gpc];
+                                        }
+                                        start = argv[2] + i + 1;
+                                }
+                                // Range start
+                                if (argv[2][i] == '-') {
+                                        argv[2][i] = '\0';
+                                        range_start_gpc = atoi(start);
+                                        start = argv[2] + i + 1;
+                                }
+                        }
+                        // Convert to string, prefix with ~, and set env var
+                        char mask_str[32+3+1]; // 32 hexits, "~0x", and '\0'
+                        snprintf(mask_str, 36, "~0x%lx%016lx", (uint64_t)(mask >> 64), (uint64_t)mask);
+                        //printf("nvtaskset: Using mask string %s\n", mask_str);
+                        setenv("LIBSMCTRL_MASK", mask_str, 1);
+                        // Start task
+                        execvp(argv[3], argv+3);
+                        error(1, errno, "Unable to launch task '%s'", argv[3]);
+                } else {
+                        // Tell libsmctrl what mask to use
+                        char* mask = malloc(strlen(argv[1]) + 2);
+                        mask[0] = '~'; // Make an enable mask
+                        strcpy(mask+1, argv[1]);
+                        setenv("LIBSMCTRL_MASK", mask, 1);
+                        free(mask); // setenv() made a copy
+                        // Start task
+                        execvp(argv[2], argv+2);
+                        error(1, errno, "Unable to launch task '%s'", argv[2]);
+                }
+        }
+        fprintf(stderr, "Invalid arguments\n");
+        return 1;
+}