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. --- .gitignore | 1 + Makefile | 23 +++-- README.md | 1 + libsmctrl.c | 163 ++++++++++++++++++++++++++++---- libsmctrl.h | 4 +- libsmctrl_test_gpc_info.c | 72 +++++++++++++-- libsmctrl_test_mask_shared.cu | 7 +- libsmctrl_test_mask_shared.h | 1 + libsmctrl_test_supreme_mask.c | 10 ++ nvtaskset.c | 210 ++++++++++++++++++++++++++++++++++++++++++ 10 files changed, 459 insertions(+), 33 deletions(-) create mode 100644 libsmctrl_test_supreme_mask.c create mode 100644 nvtaskset.c diff --git a/.gitignore b/.gitignore index 7c52826..5f0fdbe 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ libsmctrl.a libsmctrl.o libsmctrl.so libsmctrl_test_gpc_info +libsmctrl_test_supreme_mask libsmctrl_test_global_mask libsmctrl_test_stream_mask libsmctrl_test_stream_mask_override diff --git a/Makefile b/Makefile index 5256634..62ec245 100644 --- a/Makefile +++ b/Makefile @@ -19,6 +19,9 @@ libsmctrl.a: libsmctrl.c libsmctrl.h # ----- Utilities ----- # Use static linking with tests to avoid LD_LIBRARY_PATH issues +nvtaskset: nvtaskset.c libsmctrl.so libsmctrl.a + $(CC) $@.c -o $@ -L. -l:libsmctrl.a $(CFLAGS) $(LDFLAGS) + libsmctrl_test_gpc_info: libsmctrl_test_gpc_info.c libsmctrl.a testbench.h $(CC) $< -o $@ -g -L. -l:libsmctrl.a $(CFLAGS) $(LDFLAGS) @@ -26,6 +29,9 @@ libsmctrl_test_gpc_info: libsmctrl_test_gpc_info.c libsmctrl.a testbench.h libsmctrl_test_mask_shared.o: libsmctrl_test_mask_shared.cu testbench.h $(NVCC) -ccbin $(CXX) $< -c -g +libsmctrl_test_supreme_mask: libsmctrl_test_supreme_mask.c libsmctrl.a libsmctrl_test_mask_shared.o + $(NVCC) -ccbin $(CXX) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS) + libsmctrl_test_global_mask: libsmctrl_test_global_mask.c libsmctrl.a libsmctrl_test_mask_shared.o $(NVCC) -ccbin $(CXX) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS) @@ -41,12 +47,17 @@ libsmctrl_test_next_mask: libsmctrl_test_next_mask.c libsmctrl.a libsmctrl_test_ libsmctrl_test_next_mask_override: libsmctrl_test_next_mask_override.c libsmctrl.a libsmctrl_test_mask_shared.o $(NVCC) -ccbin $(CXX) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS) -tests: libsmctrl_test_gpc_info libsmctrl_test_global_mask libsmctrl_test_stream_mask libsmctrl_test_stream_mask_override libsmctrl_test_next_mask libsmctrl_test_next_mask_override +tests: libsmctrl_test_gpc_info libsmctrl_test_supreme_mask \ + libsmctrl_test_global_mask libsmctrl_test_stream_mask \ + libsmctrl_test_stream_mask_override libsmctrl_test_next_mask \ + libsmctrl_test_next_mask_override -all: libsmctrl.so tests +all: libsmctrl.so nvtaskset tests clean: - rm -f libsmctrl.so libsmctrl.a libsmctrl_test_gpu_info \ - libsmctrl_test_mask_shared.o libmsctrl_test_global_mask \ - libsmctrl_test_stream_mask libmsctrl_test_stream_mask_override \ - libsmctrl_test_next_mask libmsctrl_test_next_mask_override + rm -f libsmctrl.so libsmctrl.o libsmctrl.a libsmctrl_test_gpc_info \ + libsmctrl_test_mask_shared.o libsmctrl_test_supreme_mask \ + libsmctrl_test_global_mask \ + libsmctrl_test_stream_mask libsmctrl_test_stream_mask_override \ + libsmctrl_test_next_mask libsmctrl_test_next_mask_override \ + nvtaskset diff --git a/README.md b/README.md index f4d7cd1..f2be718 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,7 @@ make tests - Untested on non-Jetson `aarch64` platforms - Untested on CUDA 11.8, 12.0, and 12.1 on Jetson `aarch64` - Mask bit indexes do not directly correlate to software-visible TPC/SM IDs in V4 TMD/QMDs (Hopper+; compute capability 9.0). The mask bit indexes instead appear to correspond to on-chip-units, including disabled ones; i.e. the set of pre-SM-ID-remapping and pre-floorsweeping TPCs +- Tests fail when NVIDIA MPS is enabled on Volta-generation and newer GPUs, as non-physical SM IDs are returned by `%%smid` register in this configuration (see US Patent 11,307,903 by NVIDIA) ## Important Limitations diff --git a/libsmctrl.c b/libsmctrl.c index 5ee94fb..6aa471b 100644 --- a/libsmctrl.c +++ b/libsmctrl.c @@ -31,6 +31,8 @@ #include #include #include +#include +#include #include #include "libsmctrl.h" @@ -52,6 +54,8 @@ static const CUuuid callback_funcs_id = {0x2c, (char)0x8e, 0x0a, (char)0xd8, 0x0 // structures, allowing us to override it with the next mask. #define QMD_DOMAIN 0xb #define QMD_PRE_UPLOAD 0x1 +// Supreme mask (cannot be overridden) +static uint64_t *g_supreme_sm_mask = NULL; // Global mask (applies across all threads) static uint64_t g_sm_mask = 0; // Next mask (applies per-thread) @@ -110,6 +114,13 @@ static void control_callback_v2(void *ukwn, int domain, int cbid, const void *in *upper_ptr = (uint32_t)(g_sm_mask >> 32); } + // No one may override the supreme SM mask; any SMs disabled in it (set + // bits) must always remain disabled. + if (g_supreme_sm_mask) { + *lower_ptr |= (uint32_t)*g_supreme_sm_mask; + *upper_ptr |= (uint32_t)(*g_supreme_sm_mask >> 32); + } + //fprintf(stderr, "Final SM Mask (lower): %x\n", *lower_ptr); //fprintf(stderr, "Final SM Mask (upper): %x\n", *upper_ptr); } @@ -423,13 +434,30 @@ static int read_int_procfile(char* filename, uint64_t* out) { return 0; } -// We support up to 64 TPCs, up to 12 GPCs per GPU, and up to 16 GPUs. -// TODO: Handle GPUs with greater than 64 TPCs (e.g. some H100 variants) -static uint64_t tpc_mask_per_gpc_per_dev[16][12]; +// We support up to 128 TPCs, up to 12 GPCs per GPU, and up to 16 GPUs. +#define MAX_GPCS 12 +static uint64_t tpc_mask_per_gpc_per_dev[16][MAX_GPCS]; +static uint128_t tpc_mask_per_gpc_per_dev_ext[16][MAX_GPCS]; // Output mask is vtpc-indexed (virtual TPC) +// Note that this function has to undo _both_ floorsweeping and ID remapping int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc, int dev) { - uint32_t i, j, vtpc_idx = 0; - uint64_t gpc_mask, num_tpc_per_gpc, max_gpcs, gpc_tpc_mask; + int err, i; + uint128_t *tpcs_for_gpc_ext; + if ((err = libsmctrl_get_gpc_info_ext(num_enabled_gpcs, &tpcs_for_gpc_ext, dev))) + return err; + for (i = 0; i < *num_enabled_gpcs; i++) { + if ((tpcs_for_gpc_ext[i] & -1ull) != tpcs_for_gpc_ext[i]) + return ERANGE; + tpc_mask_per_gpc_per_dev[dev][i] = (uint64_t)tpcs_for_gpc_ext[i]; + } + *tpcs_for_gpc = tpc_mask_per_gpc_per_dev[dev]; + return 0; +} + +int libsmctrl_get_gpc_info_ext(uint32_t* num_enabled_gpcs, uint128_t** tpcs_for_gpc, int dev) { + uint32_t i, j, tpc_id, gpc_id, num_enabled_tpcs, num_configured_tpcs; + uint64_t gpc_mask, num_tpc_per_gpc, max_gpcs, gpc_tpc_mask, gpc_tpc_config, total_read = 0; + uint128_t tpc_bit; int err; char filename[100]; *num_enabled_gpcs = 0; @@ -449,43 +477,79 @@ int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc, snprintf(filename, 100, "/proc/gpu%d/gpc_mask", dev); if (err = read_int_procfile(filename, &gpc_mask)) return err; + // Determine the number of enabled TPCs snprintf(filename, 100, "/proc/gpu%d/num_tpc_per_gpc", dev); if (err = read_int_procfile(filename, &num_tpc_per_gpc)) return err; // For each enabled GPC + num_enabled_tpcs = 0; for (i = 0; i < max_gpcs; i++) { // Skip this GPC if disabled if ((1 << i) & gpc_mask) continue; (*num_enabled_gpcs)++; - // Get the bitstring of TPCs disabled for this GPC + // Get the bitstring of TPCs disabled for this physical GPC // Set bit = disabled TPC snprintf(filename, 100, "/proc/gpu%d/gpc%d_tpc_mask", dev, i); if (err = read_int_procfile(filename, &gpc_tpc_mask)) return err; - uint64_t* tpc_mask = &tpc_mask_per_gpc_per_dev[dev][*num_enabled_gpcs - 1]; - *tpc_mask = 0; - for (j = 0; j < num_tpc_per_gpc; j++) { - // Skip disabled TPCs - if ((1 << j) & gpc_tpc_mask) - continue; - *tpc_mask |= (1ull << vtpc_idx); - vtpc_idx++; + // Bits greater than the max number of TPCs should be ignored, so only + // keep the `num_tpc_per_gpc`-count number of lower bits. + gpc_tpc_mask &= -1u >> (64 - num_tpc_per_gpc); + // Number of enabled TPCs = max - number disabled + num_enabled_tpcs += num_tpc_per_gpc - __builtin_popcountl(gpc_tpc_mask); + } + // Clear any previous mask + for (i = 0; i < MAX_GPCS; i++) + tpc_mask_per_gpc_per_dev_ext[dev][i] = 0; + // For each enabled TPC + for (tpc_id = 0; tpc_id < num_enabled_tpcs;) { + // Pull mapping for the next set of 4 TPCs + snprintf(filename, 100, "/proc/gpu%d/CWD_GPC_TPC_ID%d", dev, tpc_id / 4); + if (err = read_int_procfile(filename, &gpc_tpc_config)) + return err; + total_read += gpc_tpc_config; + for (j = 0; j < 4 && tpc_id < num_enabled_tpcs; j++, tpc_id++) { + // Set the bit for the current TPC + tpc_bit = 1; + tpc_bit <<= tpc_id; + // Determine which GPC the current TPC is associated with + // (upper 4 bits of each byte) + gpc_id = (gpc_tpc_config >> (j*8 + 4) & 0xfu); + // Save mapping + tpc_mask_per_gpc_per_dev_ext[dev][gpc_id] |= tpc_bit; } } - *tpcs_for_gpc = tpc_mask_per_gpc_per_dev[dev]; + // Verify each TPC is configured + tpc_bit = 0; + for (i = 0; i < MAX_GPCS; i++) + tpc_bit |= tpc_mask_per_gpc_per_dev_ext[dev][i]; + num_configured_tpcs = __builtin_popcountl(tpc_bit) + __builtin_popcountl(tpc_bit >> 64); + if (num_configured_tpcs != num_enabled_tpcs) { + fprintf(stderr, "libsmctrl: Found configuration for only %d TPCs when %d were expected.\n", num_configured_tpcs, num_enabled_tpcs); + return EIO; + } + // Verify that the configuration was not always zero (indicates a powered- + // -off GPU). + if (total_read == 0) { + fprintf(stderr, "libsmctrl: Is GPU on? Configuration registers are all zero.\n"); + return EIO; + } + + *tpcs_for_gpc = tpc_mask_per_gpc_per_dev_ext[dev]; return 0; } int libsmctrl_get_tpc_info(uint32_t* num_tpcs, int dev) { uint32_t num_gpcs; - uint64_t* tpcs_per_gpc; + uint128_t* tpcs_per_gpc; int res; - if (res = libsmctrl_get_gpc_info(&num_gpcs, &tpcs_per_gpc, dev)) + if (res = libsmctrl_get_gpc_info_ext(&num_gpcs, &tpcs_per_gpc, dev)) return res; *num_tpcs = 0; for (int gpc = 0; gpc < num_gpcs; gpc++) { *num_tpcs += __builtin_popcountl(tpcs_per_gpc[gpc]); + *num_tpcs += __builtin_popcountl(tpcs_per_gpc[gpc] >> 64); } return 0; } @@ -526,3 +590,68 @@ abort_cuda: return EIO; } +// Allow setting a default mask via an environment variable +// Also enables libsmctrl to be used on unmodified programs via setting: +// LD_PRELOAD=libsmctrl.so LIBSMCTRL_MASK= ./my_program +// Where "" is replaced with a disable mask, optionally prefixed +// with a ~ to invert it (make it an enable mask). +__attribute__((constructor)) static void setup(void) { + char *end, *mask_str; + // If dynamic changes are disabled (due to an error) this variable is + // permanently used to store the supreme mask, rather than the SysV shared + // memory segment. + static uint64_t mask; + bool invert = false; + int shmid; + key_t shm_key; + + mask_str = getenv("LIBSMCTRL_MASK"); + if (!mask_str) + return; + + if (*mask_str == '~') { + invert = true; + mask_str++; + } + + // XXX: Doesn't support 128-bit masks + mask = strtoull(mask_str, &end, 0); + // Verify we were able to parse the whole string + if (*end != '\0') + abort(1, EINVAL, "Unable to apply default mask"); + + if (invert) + mask = ~mask; + + // Initialize CUDA and the interception callback + setup_sm_control_callback(); + + // TODO: Switch to memfd_create(); this leaks IPC objects + // Create a SysV IPC key (32 bits) to identify our shared memory region + // Use the pid as the top 16 bits, and "sm" as the bottom 16 + shm_key = getpid(); + shm_key <<= 16; + shm_key |= (int)'s' << 8 | (int)'m'; + // Obtain or create a 128-bit (16-byte) shared memory region + shmid = shmget(shm_key, 16, IPC_CREAT | 0600); + if (shmid == -1) { + abort(0, errno, "Unable to create shared memory for dynamic partition changes. Dynamic changes disabled"); + g_supreme_sm_mask = &mask; + return; + } + // Open the shared memory region + g_supreme_sm_mask = shmat(shmid, NULL, 0); + if (g_supreme_sm_mask == (void*)-1) { + abort(0, errno, "Unable to create shared memory for dynamic partition changes. Dynamic changes disabled"); + g_supreme_sm_mask = &mask; + return; + } + // XXX: This makes the region unopenable to everyone else. Switch to memfd! + // Mark the shared memory region for deletion (after we terminate) + if (shmctl(shmid, IPC_RMID, NULL) == -1) + abort(0, errno, "Unable to mark shared memory for dynamic partition changes for deletion on process termination. Will leak one page of memory."); + + // Set the super-global mask which cannot be overwritten by any libsmctrl + // API function. + *g_supreme_sm_mask = mask; +} diff --git a/libsmctrl.h b/libsmctrl.h index b85c0c7..45d2cd7 100644 --- a/libsmctrl.h +++ b/libsmctrl.h @@ -1,5 +1,5 @@ /** - * Copyright 2022-2024 Joshua Bakita + * Copyright 2022-2025 Joshua Bakita * Library to control TPC masks on CUDA launches. Co-opts preexisting debug * logic in the CUDA driver library, and thus requires a build with -lcuda. */ @@ -58,11 +58,13 @@ extern void libsmctrl_set_next_mask(uint64_t mask); // Get number of GPCs for devices number `dev`, and a GPC-indexed array // containing masks of which TPCs are associated with each GPC. // Note that the `nvdebug` module must be loaded to use this function. +// Note that at least one GPU context must be active to use this function. // @param num_enabled_gpcs (out) Location to store number of GPCs in // @param tpcs_for_gpc (out) Pointer to store pointer to output buffer at // @param dev (in) `nvdebug` device ID // @return 0 on success, `errno`-compatible error code on failure extern int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc, int dev); +extern int libsmctrl_get_gpc_info_ext(uint32_t* num_enabled_gpcs, uint128_t** tpcs_for_gpc, int dev); // Get total number of TPCs on device number `dev`. Requires `nvdebug`. // @param num_tpcs (out) Location to store number of TPCs at // @param dev (in) `nvdebug` device ID diff --git a/libsmctrl_test_gpc_info.c b/libsmctrl_test_gpc_info.c index 9b2b963..558b80a 100644 --- a/libsmctrl_test_gpc_info.c +++ b/libsmctrl_test_gpc_info.c @@ -1,7 +1,18 @@ -// Copyright 2024 Joshua Bakita +// Copyright 2024-2025 Joshua Bakita +// Obtain and print the correspondence between TPCs and GPCs for a given GPU. +// +// Known issues: +// - If CUDA cannot see the same number of GPUs as the nvdebug kernel module, +// the passed GPU ID may not properly correspond to to an ID an CUDA. This +// will cause us to fail to initialize a context on the right device, and +// may cause the test to terminate due to no initialized context. This should +// only happen if some of the attached GPUs are too old or new for CUDA. #define _GNU_SOURCE +#include + #include #include +#include #include #include #include @@ -11,9 +22,9 @@ int main(int argc, char** argv) { uint32_t num_gpcs = 0, num_tpcs = 0; - uint64_t* masks = NULL; - int res; - int gpu_id = 0; + uint128_t* masks = NULL; + int res, print_width, gpu_id; + CUcontext ctx; // Optionally support specifying the GPU ID to query via an argument // Important: This GPU ID must match the ID used by the nvdebug module. See // the documentation on libsmctrl_get_gpc_info() for details. @@ -23,12 +34,57 @@ int main(int argc, char** argv) { } if (argc > 1) gpu_id = atoi(argv[1]); - if ((res = libsmctrl_get_gpc_info(&num_gpcs, &masks, gpu_id)) != 0) - error(1, res, "libsmctrl_get_gpc_info() failed"); + else + gpu_id = 0; + // 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); + return 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); + return 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); + return 1; + } printf("%s: GPU%d has %d enabled GPCs.\n", program_invocation_name, gpu_id, num_gpcs); + // Determine how wide the print should be (for pretty-printing) + print_width = 0; + for (int i = 0; i < num_gpcs; i++) { + int shift = 0; + while (masks[i] >> shift) + shift++; + if (shift > print_width) + print_width = shift; + } + // Convert the width to a number of octets, rather than number of bits + // (Result of integer divison, +1 if it does not evenly divide) + print_width = print_width/4 + !!(print_width % 4); for (int i = 0; i < num_gpcs; i++) { - num_tpcs += __builtin_popcountl(masks[i]); - printf("%s: Mask of %d TPCs associated with GPC %d: %#018lx\n", program_invocation_name, __builtin_popcountl(masks[i]), i, masks[i]); + // No built-in for 128-bit integers, so split it into two 64-bit ones + int num_tpcs_local = __builtin_popcountl(masks[i]) + __builtin_popcountl(masks[i] >> 64); + num_tpcs += num_tpcs_local; + if (print_width > 16) + printf("%s: Mask of %d TPCs associated with GPC %d: 0x%0*lx%016lx\n", + program_invocation_name, num_tpcs_local, i, print_width - 16, + (uint64_t)(masks[i] >> 64), (uint64_t)masks[i]); + else + printf("%s: Mask of %d TPCs associated with GPC %d: 0x%0*lx\n", + program_invocation_name, num_tpcs_local, i, print_width, + (uint64_t)masks[i]); } printf("%s: Total of %u enabled TPCs.\n", program_invocation_name, num_tpcs); return 0; diff --git a/libsmctrl_test_mask_shared.cu b/libsmctrl_test_mask_shared.cu index e4ddde7..3b7ebcd 100644 --- a/libsmctrl_test_mask_shared.cu +++ b/libsmctrl_test_mask_shared.cu @@ -1,9 +1,10 @@ -// Copyright 2023 Joshua Bakita +// Copyright 2023-2025 Joshua Bakita #include #include #include #include #include +#include // For getpid() #include "libsmctrl.h" #include "testbench.h" @@ -81,6 +82,10 @@ int test_constrained_size_and_location(enum partitioning_type part_type) { // Apply partitioning to enable only the first TPC of each 32-bit block switch (part_type) { + case PARTITION_SUPREME: + printf("%s: Please set mask to '0x%016lx%016lx' for PID %d using the control deamon and press any key to continue...\n", program_invocation_name, (uint64_t)(mask >> 64), (uint64_t)mask, getpid()); + fgetc(stdin); + break; case PARTITION_GLOBAL: libsmctrl_set_global_mask(mask); break; diff --git a/libsmctrl_test_mask_shared.h b/libsmctrl_test_mask_shared.h index 5a115b7..14ee2e7 100644 --- a/libsmctrl_test_mask_shared.h +++ b/libsmctrl_test_mask_shared.h @@ -4,6 +4,7 @@ extern "C" { #endif enum partitioning_type { + PARTITION_SUPREME, PARTITION_GLOBAL, PARTITION_STREAM, PARTITION_STREAM_OVERRIDE, diff --git a/libsmctrl_test_supreme_mask.c b/libsmctrl_test_supreme_mask.c new file mode 100644 index 0000000..620f5cf --- /dev/null +++ b/libsmctrl_test_supreme_mask.c @@ -0,0 +1,10 @@ +// Copyright 2025 Joshua Bakita +// Test environment variable LIBSMCTRL_MASK +// All types of partitioning use the same test, so this file is trival. + +#include "libsmctrl_test_mask_shared.h" + +int main() { + return test_constrained_size_and_location(PARTITION_SUPREME); +} + 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