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 are limited to 108 characters long
+        snprintf(mps_ctrl_addr.sun_path, 108, "%s/control", mps_pipe_dir);
+        // This mirrors the process `nvidia-cuda-mps-control` uses to detect MPS
+        if ((mps_ctrl = socket(AF_UNIX, SOCK_SEQPACKET, 0)) == -1)
+                return false;
+        if (setsockopt(mps_ctrl, SOL_SOCKET, SO_PASSCRED, &yes, sizeof(yes)) == -1)
+                return false;
+        if (connect(mps_ctrl, &mps_ctrl_addr, sizeof(struct sockaddr_un)) == -1)
+                return false;