From 97e97dc9d55e71edbc2031e6a509a7cc17abe168 Mon Sep 17 00:00:00 2001 From: Joshua Bakita Date: Mon, 5 May 2025 03:30:13 -0400 Subject: Major update for ECRTS'25: fix TPC to GPU mapping and add a "supreme" mask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These updates are featured in the paper: J. Bakita and J. H. Anderson, “Hardware Compute Partitioning on NVIDIA GPUs for Composable Systems”, Proceedings of the 37th Euromicro Conference on Real-Time Systems (ECRTS), to appear, Jul 2025. They: 1. Fix reported GPC to TPC mappings (requires nvdebug update). 2. Add support for a "supreme" mask, which overrides all others and can be set on a per-process basis via an environment variable, and optionally modified at runtime via the nvtaskset utility. 3. Add test for the supreme mask. --- nvtaskset.c | 210 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 nvtaskset.c (limited to 'nvtaskset.c') 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include // To help with getting GPC info + +#include "libsmctrl.h" + +const char* maintainer = ""; +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 \n", argv[0]); + fprintf(stderr, " %s \n", argv[0]); + fprintf(stderr, " %s --gpc-list \n", argv[0]); + fprintf(stderr, " 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; +} -- cgit v1.2.2