1 files changed, 26 insertions, 23 deletions

diff --git a/libsmctrl.c b/libsmctrl.c
index 09fc627..5ee94fb 100644
--- a/libsmctrl.c
+++ b/libsmctrl.c
@@ -1,5 +1,5 @@
 /**
- * Copyright 2022-2024 Joshua Bakita
+ * Copyright 2022-2025 Joshua Bakita
  * Library to control SM masks on CUDA launches. Co-opts preexisting debug
  * logic in the CUDA driver library, and thus requires a build with -lcuda.
  *
@@ -10,7 +10,7 @@
  *   +-----------+---------------+---------------+--------------+
  *   |  Version  |  Global Mask  |  Stream Mask  |  Next Mask   |
  *   +-----------+---------------+---------------+--------------+
- *   | 8.0-12.6  | TMD/QMD Hook  | stream struct | TMD/QMD Hook |
+ *   | 8.0-12.8  | TMD/QMD Hook  | stream struct | TMD/QMD Hook |
  *   | 6.5-7.5   | TMD/QMD Hook  | N/A           | TMD/QMD Hook |
  *   +-----------+---------------+---------------+--------------+
  * "N/A" indicates that a mask type is unsupported on that CUDA version.
@@ -119,10 +119,23 @@ static void setup_sm_control_callback() {
         int (*enable)(uint32_t enable, uint32_t hndl, int domain, int cbid);
         uintptr_t* tbl_base;
         uint32_t my_hndl;
-        // Avoid race conditions (setup can only be called once)
+        // Avoid race conditions (setup should only run once)
         if (__atomic_test_and_set(&sm_control_setup_called, __ATOMIC_SEQ_CST))
                 return;
 
+#if CUDA_VERSION <= 6050
+        // Verify supported CUDA version
+        // It's impossible for us to run with a version of CUDA older than we were
+        // built by, so this check is excluded if built with CUDA > 6.5.
+        int ver = 0;
+        cuDriverGetVersion(&ver);
+        if (ver < 6050)
+                abort(1, ENOSYS, "Global or next masking requires at least CUDA 6.5; "
+                                 "this application is using CUDA %d.%d",
+                                 ver / 1000, (ver % 100));
+#endif
+
+        // Set up callback
         cuGetExportTable((const void**)&tbl_base, &callback_funcs_id);
         uintptr_t subscribe_func_addr = *(tbl_base + 3);
         uintptr_t enable_func_addr = *(tbl_base + 6);
@@ -139,31 +152,13 @@ static void setup_sm_control_callback() {
 
 // Set default mask for all launches
 void libsmctrl_set_global_mask(uint64_t mask) {
-        if (!sm_control_setup_called) {
-                // The version will not change while running, so only check once
-                int ver = 0;
-                cuDriverGetVersion(&ver);
-                if (ver < 6050)
-                        abort(1, ENOSYS, "Global masking requires at least CUDA 6.5; "
-                                         "this application is using CUDA %d.%d",
-                                         ver / 1000, (ver % 100));
-                setup_sm_control_callback();
-        }
+        setup_sm_control_callback();
         g_sm_mask = mask;
 }
 
 // Set mask for next launch from this thread
 void libsmctrl_set_next_mask(uint64_t mask) {
-        if (!sm_control_setup_called) {
-                // The version will not change while running, so only check once
-                int ver = 0;
-                cuDriverGetVersion(&ver);
-                if (ver < 6050)
-                        abort(1, ENOSYS, "Next masking requires at least CUDA 6.5; "
-                                         "this application is using CUDA %d.%d",
-                                         ver / 1000, (ver % 100));
-                setup_sm_control_callback();
-        }
+        setup_sm_control_callback();
         g_next_sm_mask = mask;
 }
 
@@ -204,6 +199,10 @@ void libsmctrl_set_next_mask(uint64_t mask) {
 // CUDA 12.5 and 12.6 use the same offset
 // 12.5 tested on 555.58.02
 // 12.6 tested on 560.35.03
+#define CU_12_7_MASK_OFF 0x4fc
+// CUDA 12.7 and 12.8 use the same offset
+// 12.7 tested on 565.77
+// 12.8 tested on 570.124.06
 
 // Offsets for the stream struct on Jetson aarch64
 #define CU_9_0_MASK_OFF_JETSON 0x128
@@ -334,6 +333,10 @@ void libsmctrl_set_stream_mask_ext(void* stream, uint128_t mask) {
         case 12060:
                 hw_mask_v2 = (void*)(stream_struct_base + CU_12_5_MASK_OFF);
                 break;
+        case 12070:
+        case 12080:
+                hw_mask_v2 = (void*)(stream_struct_base + CU_12_7_MASK_OFF);
+                break;
 #elif __aarch64__
         case 9000: {
                 // Jetson TX2 offset is slightly different on CUDA 9.0.