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, 173 insertions, 39 deletions

diff --git a/libsmctrl.c b/libsmctrl.c
index 6aa471b..79d2b33 100644
--- a/libsmctrl.c
+++ b/libsmctrl.c
@@ -17,22 +17,27 @@
  * Please contact the authors if support is needed for a particular feature on
  * an older CUDA version. Support for those is unimplemented, not impossible.
  *
- * An old implementation of this file effected the global mask on CUDA 10.2 by
+ * An old implementation of this file affected the global mask on CUDA 10.2 by
  * changing a field in CUDA's global struct that CUDA applies to the QMD/TMD.
  * That implementation was extraordinarily complicated, and was replaced in
  * 2024 with a more-backward-compatible way of hooking the TMD/QMD.
  * View the old implementation via Git: `git show aa63a02e:libsmctrl.c`.
  */
+#define _GNU_SOURCE // To enable use of memfd_create()
 #include <cuda.h>
 
 #include <errno.h>
 #include <error.h>
+#include <dlfcn.h>
 #include <fcntl.h>
 #include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
-#include <sys/ipc.h>
-#include <sys/shm.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <sys/socket.h>
+#include <sys/types.h>
+#include <sys/un.h>
 #include <unistd.h>
 
 #include "libsmctrl.h"
@@ -48,20 +53,33 @@
 // (No testing attempted on pre-CUDA-6.5 versions)
 // Values for the following three lines can be extracted by tracing CUPTI as
 // it interects with libcuda.so to set callbacks.
-static const CUuuid callback_funcs_id = {0x2c, (char)0x8e, 0x0a, (char)0xd8, 0x07, 0x10, (char)0xab, 0x4e, (char)0x90, (char)0xdd, 0x54, 0x71, (char)0x9f, (char)0xe5, (char)0xf7, 0x4b};
+static const CUuuid callback_funcs_id = {{0x2c, (char)0x8e, 0x0a, (char)0xd8, 0x07, 0x10, (char)0xab, 0x4e, (char)0x90, (char)0xdd, 0x54, 0x71, (char)0x9f, (char)0xe5, (char)0xf7, 0x4b}};
 // These callback descriptors appear to intercept the TMD/QMD late enough that
 // CUDA has already applied the per-stream mask from its internal data
 // structures, allowing us to override it with the next mask.
 #define QMD_DOMAIN 0xb
 #define QMD_PRE_UPLOAD 0x1
+/**
+ * These globals must be non-static (i.e., have global linkage) to ensure that
+ * if multiple copies of the library are loaded (e.g., dynamically linked to
+ * both this program and a dependency), secondary copies do not attempt to
+ * repeat initialization or make changes to unused copies of mask values.
+ */
 // Supreme mask (cannot be overridden)
-static uint64_t *g_supreme_sm_mask = NULL;
+uint128_t *g_supreme_sm_mask = NULL;
 // Global mask (applies across all threads)
-static uint64_t g_sm_mask = 0;
+uint64_t g_sm_mask = 0;
 // Next mask (applies per-thread)
-static __thread uint64_t g_next_sm_mask = 0;
+__thread uint64_t g_next_sm_mask = 0;
 // Flag value to indicate if setup has been completed
-static bool sm_control_setup_called = false;
+bool sm_control_setup_called = false;
+
+#ifdef LIBSMCTRL_STATIC
+// Special handling for if built as a static library, and the libcuda.so.1
+// libsmctrl wrapper is in use (see comment on setup() constructor for detail).
+static void (*shared_set_global_mask)(uint64_t) = NULL;
+static void (*shared_set_next_mask)(uint64_t) = NULL;
+#endif
 
 // v1 has been removed---it intercepted the TMD/QMD too early, making it
 // impossible to override the CUDA-injected stream mask with the next mask.
@@ -78,7 +96,7 @@ static void control_callback_v2(void *ukwn, int domain, int cbid, const void *in
         if (!tmd)
                 abort(1, 0, "TMD allocation appears NULL; likely forward-compatibilty issue.\n");
 
-        uint32_t *lower_ptr, *upper_ptr;
+        uint32_t *lower_ptr, *upper_ptr, *ext_lower_ptr, *ext_upper_ptr;
 
         // The location of the TMD version field seems consistent across versions
         uint8_t tmd_ver = *(uint8_t*)(tmd + 72);
@@ -87,10 +105,12 @@ static void control_callback_v2(void *ukwn, int domain, int cbid, const void *in
                 // TMD V04_00 is used starting with Hopper to support masking >64 TPCs
                 lower_ptr = tmd + 304;
                 upper_ptr = tmd + 308;
+                ext_lower_ptr = tmd + 312;
+                ext_upper_ptr = tmd + 316;
                 // XXX: Disable upper 64 TPCs until we have ...next_mask_ext and
                 //      ...global_mask_ext
-                *(uint32_t*)(tmd + 312) = -1;
-                *(uint32_t*)(tmd + 316) = -1;
+                *ext_lower_ptr = -1;
+                *ext_upper_ptr = -1;
                 // An enable bit is also required
                 *(uint32_t*)tmd |= 0x80000000;
         } else if (tmd_ver >= 0x16) {
@@ -119,6 +139,10 @@ static void control_callback_v2(void *ukwn, int domain, int cbid, const void *in
         if (g_supreme_sm_mask) {
                 *lower_ptr |= (uint32_t)*g_supreme_sm_mask;
                 *upper_ptr |= (uint32_t)(*g_supreme_sm_mask >> 32);
+                if (tmd_ver >= 0x40) {
+                        *ext_lower_ptr |= (uint32_t)(*g_supreme_sm_mask >> 64);
+                        *ext_upper_ptr |= (uint32_t)(*g_supreme_sm_mask >> 96);
+                }
         }
 
         //fprintf(stderr, "Final SM Mask (lower): %x\n", *lower_ptr);
@@ -163,12 +187,26 @@ static void setup_sm_control_callback() {
 
 // Set default mask for all launches
 void libsmctrl_set_global_mask(uint64_t mask) {
+#ifdef LIBSMCTRL_STATIC
+        // Special handling for if built as a static library, and the libcuda.so.1
+        // libsmctrl wrapper is in use (see comment on setup() constructor for
+        // detail).
+        if (shared_set_global_mask)
+                return (*shared_set_global_mask)(mask);
+#endif
         setup_sm_control_callback();
         g_sm_mask = mask;
 }
 
 // Set mask for next launch from this thread
 void libsmctrl_set_next_mask(uint64_t mask) {
+#ifdef LIBSMCTRL_STATIC
+        // Special handling for if built as a static library, and the libcuda.so.1
+        // libsmctrl wrapper is in use (see comment on setup() constructor for
+        // detail).
+        if (shared_set_next_mask)
+                return (*shared_set_next_mask)(mask);
+#endif
         setup_sm_control_callback();
         g_next_sm_mask = mask;
 }
@@ -248,7 +286,7 @@ struct stream_sm_mask_v2 {
 // (CUDA 9.0 behaves slightly different on this platform.)
 // @return 1 if detected, 0 if not, -cuda_err on error
 #if __aarch64__
-int detect_parker_soc() {
+static int detect_parker_soc() {
         int cap_major, cap_minor, err, dev_count;
         if (err = cuDeviceGetCount(&dev_count))
                 return -err;
@@ -272,7 +310,7 @@ int detect_parker_soc() {
 }
 #endif // __aarch64__
 
-// Should work for CUDA 8.0 through 12.6
+// Should work for CUDA 8.0 through 12.8
 // A cudaStream_t is a CUstream*. We use void* to avoid a cuda.h dependency in
 // our header
 void libsmctrl_set_stream_mask(void* stream, uint64_t mask) {
@@ -417,7 +455,8 @@ void libsmctrl_set_stream_mask_ext(void* stream, uint128_t mask) {
         }
 }
 
-/* INFORMATIONAL FUNCTIONS */
+
+/*** TPC and GPU Informational Functions ***/
 
 // Read an integer from a file in `/proc`
 static int read_int_procfile(char* filename, uint64_t* out) {
@@ -590,32 +629,98 @@ abort_cuda:
         return EIO;
 }
 
+
+/*** Private functions for nvtaskset and building as a libcuda.so.1 wrapper ***/
+
+// Check if NVIDIA MPS is running, following the process that `strace` shows
+// `nvidia-cuda-mps-control` to use. MPS is a prerequisite to co-running
+// multiple GPU-using tasks without timeslicing.
+bool libsmctrl_is_mps_running() {
+        char *mps_pipe_dir;
+        int mps_ctrl;
+        struct sockaddr_un mps_ctrl_addr;
+        mps_ctrl_addr.sun_family = AF_UNIX;
+        const int yes = 1;
+
+        if (!(mps_pipe_dir = getenv("CUDA_MPS_PIPE_DIRECTORY")))
+                mps_pipe_dir = "/tmp/nvidia-mps";
+        // Pipe names