aboutsummaryrefslogtreecommitdiffstats
path: root/nvtaskset.c
diff options
context:
space:
mode:
Diffstat (limited to 'nvtaskset.c')
-rw-r--r--nvtaskset.c520
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 @@
1// Copyright 2025 Joshua Bakita 1// Copyright 2025 Joshua Bakita
2// taskset-like utility for the GPU 2// Show or change the GPU core affinity for a CUDA process
3// taskset-like utility for NVIDIA GPUs
3#define _GNU_SOURCE // For program_invocation_name 4#define _GNU_SOURCE // For program_invocation_name
4#include <argp.h> 5#include <argp.h>
6#include <dirent.h>
5#include <errno.h> 7#include <errno.h>
6#include <error.h> 8#include <error.h>
9#include <fcntl.h>
10#include <stdbool.h>
7#include <stdint.h> 11#include <stdint.h>
8#include <stdio.h> 12#include <stdio.h>
9#include <stdlib.h> 13#include <stdlib.h>
10#include <string.h> 14#include <string.h>
11#include <sys/ipc.h> 15#include <sys/mman.h>
12#include <sys/shm.h>
13#include <sys/types.h>
14#include <unistd.h> 16#include <unistd.h>
15 17
16#include <cuda.h> // To help with getting GPC info 18#include <cuda.h> // To help with getting GPC info
17 19
18#include "libsmctrl.h" 20#include "libsmctrl.h"
19 21
20const char* maintainer = "<jbakita@cs.unc.edu>"; 22#define LINK_NAME "/memfd:libsmctrl"
21const char* version = "nvtaskset 2025.03"; 23
22const char* desc = "taskset-like utility for NVIDIA GPUs."; 24// TODO: Write automated tests:
25// - Change region of non-existent PID
26// - Change region of permission denied PID
27// - Change region of non-GPU PID
28// - Change GPC list
29// - Change TPC list
30// - Change TPC mask
31// - Start TPC mask
32// - Start TPC list
33// - Start GPC list
34// - Start with subargument containing -
35// - Query GPC list
36// - Query TPC list
37// - Query TPC mask
38// - Set GPC list w/ non-existant GPC
39// - Set TPC list w/ non-existant TPC
40
41// Private symbols from libsmctrl
42extern bool libsmctrl_is_mps_running();
43extern uint128_t strtou128(const char *nptr, char **endptr, int base);
44
45const char *argp_program_bug_address = "<jbakita@cs.unc.edu>";
46const char *argp_program_version = "nvtaskset 2025.06";
47const char desc[] = "Show or change the GPU core affinity for a CUDA process\v"
48 "Warning: When using GPC lists, this tool currently "
49 "derives TPC to GPC mappings from the first NVIDIA GPU in "
50 "the system (by PCI bus ID) device. To use the mappings "
51 "for a different device, use the `libsmctrl_test_get_info` "
52 "tool to get the bitmask of TPCs associated with each GPC, "
53 "OR them, and then set that bitmask via this tool. Better "
54 "multi-GPU support is intended for a future release.\n\n"
55 "Inspired by the Linux taskset utility.";
56const char args_doc[] = "[mask | list] [pid | cmd [args...]]";
23 57
24const struct argp_option opts[] = { 58const struct argp_option opts[] = {
59 {"gpc-list", 'g', NULL, 0, "Specify partition as a list of GPCs"},
60 {"tpc-list", 't', NULL, 0, "Specify partition as a list of TPCs"},
61 {"pid", 'p', NULL, 0, "Operate on an existing PID"},
25 {0} 62 {0}
26}; 63};
27 64
28unsigned __int128 strtou128(const char *nptr, char **endptr, int base) { 65// Create a CUDA context and query the associated GPC to TPC mappings
29 unsigned __int128 result = 0; 66// Based off logic in libsmctrl_test_gpc_info
30 if (base != 16)
31 error(1, EINVAL, "Internal error");
32 // Skip a "0x" prefix. Safe due to early evaluation
33 if (*nptr == '0' && (*(nptr + 1) == 'x' || *(nptr + 1) == 'X'))
34 nptr += 2;
35 // Until hitting an invalid character
36 while (1) {
37 if (*nptr >= 'a' && *nptr <= 'f')
38 result = result << 4 | (*nptr - 'a' + 10);
39 else if (*nptr >= 'A' && *nptr <= 'F')
40 result = result << 4 | (*nptr - 'A' + 10);
41 else if (*nptr >= '0' && *nptr <= '9')
42 result = result << 4 | (*nptr - '0');
43 else
44 break;
45 nptr++;
46 }
47 if (endptr)
48 *endptr = (char*)nptr;
49 return result;
50}
51
52void libsmctrl_get_gpc_info_ext_easy(uint32_t* num_gpcs, uint128_t** masks, int gpu_id) { 67void libsmctrl_get_gpc_info_ext_easy(uint32_t* num_gpcs, uint128_t** masks, int gpu_id) {
53 int res; 68 int res;
54 CUcontext ctx; 69 CUcontext ctx;
55 // XXX: Copied from libsmctrl_test_gpc_info 70 char *old_order = NULL;
56 // Tell CUDA to use PCI device id ordering (to match nvdebug) 71 // Tell CUDA to use PCI device id ordering (to match nvdebug)
57 putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID"); 72 putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID");
58 // A CUDA context is required before reading the topology information 73 // Allow CUDA to see all devices (to better match nvdebug)
59 if ((res = cuInit(0))) { 74 if (getenv("CUDA_VISIBLE_DEVICES")) {
60 const char* name; 75 if (!(old_order = strdup(getenv("CUDA_VISIBLE_DEVICES"))))
61 cuGetErrorName(res, &name); 76 error(1, errno, "Unable to allocate environment string");
62 fprintf(stderr, "%s: Unable to initialize CUDA, error %s\n", program_invocation_name, name); 77 unsetenv("CUDA_VISIBLE_DEVICES");
63 exit(1); 78 }
64 } 79 // A CUDA context is required before reading the topology information
65 if ((res = cuCtxCreate(&ctx, 0, 0))) { 80 if ((res = cuInit(0))) {
66 const char* name; 81 const char* name;
67 cuGetErrorName(res, &name); 82 cuGetErrorName(res, &name);
68 fprintf(stderr, "%s: Unable to create a CUDA context, error %s\n", program_invocation_name, name); 83 error(1, 0, "Unable to create a initialize CUDA, error %s\n", name);
69 exit(1); 84 }
70 } 85 if ((res = cuCtxCreate(&ctx, 0, gpu_id))) {
71 // Pull topology information from libsmctrl 86 const char* name;
72 if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) { 87 cuGetErrorName(res, &name);
73 error(0, res, "libsmctrl_get_gpc_info() failed"); 88 error(1, 0, "Unable to create a CUDA context, error %s\n", name);
74 if (res == ENOENT) 89 }
75 fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name); 90 // Pull topology information from libsmctrl
76 if (res == EIO) 91 if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) {
77 fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name); 92 error(0, res, "libsmctrl_get_gpc_info() failed");
78 exit(1); 93 if (res == ENOENT)
79 } 94 fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name);
80 // Not copied 95 if (res == EIO)
96 fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name);
97 exit(1);
98 }
99 // Restore the environment (in case we exec() later)
81 unsetenv("CUDA_DEVICE_ORDER"); 100 unsetenv("CUDA_DEVICE_ORDER");
101 if (old_order) {
102 setenv("CUDA_VISIBLE_DEVICES", old_order, 1);
103 free(old_order);
104 }
82} 105}
83 106
84int main(int argc, char **argv) { 107int parse_list(bool use_gpcs, char* list, uint128_t *mask_out) {
85 if (argc < 3) { 108 // We support the same ranges as taskset, e.g., X,Y,Z and X,Y-Z
86 fprintf(stderr, "Usage: %s -p <hex mask> <pid>\n", argv[0]); 109 uint32_t num_xpcs = 0; // Either TPC or GPC count, i.e., "X"PC
87 fprintf(stderr, " %s <hex mask> <command> <argument...>\n", argv[0]); 110 uint128_t* masks = NULL;
88 fprintf(stderr, " %s --gpc-list <gpc list> <command> <argument...>\n", argv[0]); 111 // TODO: Allow specifying GPU ID, rather than assuming 0!
89 fprintf(stderr, " <hex mask> has a bit set for each TPC to be enabled\n"); 112 if (use_gpcs)
90 return 1; 113 libsmctrl_get_gpc_info_ext_easy(&num_xpcs, &masks, 0);
91 } 114 else
92 // TODO: Use a proper argument parser 115 libsmctrl_get_tpc_info_cuda(&num_xpcs, 0);
93 if (strcmp("-p", argv[1]) == 0) { // Setting mask on running task 116 uint128_t mask = 0;
94 char *end; 117 int range_start_xpc = -1;
95 pid_t target_pid = strtoul(argv[2], &end, 10); 118 char* start = list;
96 // strtoul stores a pointer to the first invalid character in `end` 119 int len = strlen(list);
97 if (*end != '\0') { 120 // Convert comma-seperated GPC/TPC list into a mask
98 fprintf(stderr, "Invalid character \"%c\" in PID argument.\n", *end); 121 for (int i = 0; i < len + 1; i++) {
99 return 1; 122 if (list[i] == ',' || list[i] == '\0') {
123 list[i] = '\0';
124 int xpc = atoi(start);
125 if (xpc > num_xpcs - 1)
126 error(1, EINVAL, "%s is not a valid %s ID", start, use_gpcs ? "GPC" : "TPC");
127 // Handle ranges
128 if (range_start_xpc != -1) {
129 if (range_start_xpc >= xpc)
130 error(1, EINVAL, "Malformed %s range", use_gpcs ? "GPC" : "TPC");
131 while (range_start_xpc <= xpc) {
132 if (use_gpcs)
133 mask |= masks[range_start_xpc];
134 else
135 mask |= (uint128_t)1 << range_start_xpc;
136 range_start_xpc++;
137 }
138 range_start_xpc = -1;
139 } else {
140 if (use_gpcs)
141 mask |= masks[xpc];
142 else
143 mask |= (uint128_t)1 << xpc;
144 }
145 start = list + i + 1;
100 } 146 }
101 unsigned __int128 mask = strtou128(argv[3], &end, 16); 147 // Range start
102 if (*end != '\0') { 148 if (list[i] == '-') {
103 fprintf(stderr, "Invalid character \"%c\" in mask argument.\n", *end); 149 list[i] = '\0';
104 return 1; 150 range_start_xpc = atoi(start);
151 start = list + i + 1;
105 } 152 }
106 // The shared memory lookup key is the lower 16-bits of the PID | "sm" 153 }
107 key_t shm_key = target_pid << 16 | (int)'s' << 8 | (int) 'm'; 154 *mask_out = mask;
108