10 files changed, 459 insertions, 33 deletions

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 <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
+#include <sys/ipc.h>
+#include <sys/shm.h>
 #include <unistd.h>
 
 #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_i