Support stream masking on CUDA 12.4 (x86) and 12.6 (x86, aarch64)

Credit to Nordine Feddal for testing CUDA 12.4 on 550.544.14.

3 files changed, 26 insertions, 5 deletions

diff --git a/README.md b/README.md
index 3689ecd..0c99113 100644
--- a/README.md
+++ b/README.md
@@ -93,7 +93,7 @@ make tests
 #### Known Working
 
 - NVIDIA GPUs from compute capability 3.5 through 8.9, including embedded "Jetson" GPUs
-- CUDA 8.1 through 12.2
+- CUDA 8.0 through 12.2, plus 12.4 and 12.6
 - `x86_64` and Jetson `aarch64` platforms
 
 #### Known Issues
diff --git a/libsmctrl.c b/libsmctrl.c
index 817cb5d..b10b885 100644
--- a/libsmctrl.c
+++ b/libsmctrl.c
@@ -11,7 +11,7 @@
  *   +-----------+---------------+---------------+--------------+
  *   |  Version  |  Global Mask  |  Stream Mask  |  Next Mask   |
  *   +-----------+---------------+---------------+--------------+
- *   | 11.0-12.2 | TMD/QMD Hook  | stream struct | TMD/QMD Hook |
+ *   | 11.0-12.6 | TMD/QMD Hook  | stream struct | TMD/QMD Hook |
  *   | 10.2      | global struct | stream struct | N/A          |
  *   | 8.0-10.1  | N/A           | stream struct | N/A          |
  *   +-----------+---------------+---------------+--------------+
@@ -269,12 +269,19 @@ void libsmctrl_set_next_mask(uint64_t mask) {
 // 12.0 tested on 525.147.05
 #define CU_12_2_MASK_OFF 0x4e4
 // 12.2 tested on 535.129.03
+// CUDA 12.3 UNTESTED
+#define CU_12_4_MASK_OFF 0x4ac
+// 12.4 tested on 550.54.14 and 550.54.15
+// CUDA 12.5 UNTESTED
+#define CU_12_6_MASK_OFF 0x4ec
+// 12.6 tested on 560.35.03
 
 // Offsets for the stream struct on aarch64
 // All tested on Nov 13th, 2023
 #define CU_9_0_MASK_OFF_JETSON 0x128 // Tested on TX2
 #define CU_10_2_MASK_OFF_JETSON 0x24c // Tested on TX2 and Jetson Xavier
 #define CU_11_4_MASK_OFF_JETSON 0x394 // Tested on Jetson Orin
+#define CU_12_6_MASK_OFF_JETSON 0x514 // Tested on Jetson Orin
 
 // Used up through CUDA 11.8 in the stream struct
 struct stream_sm_mask {
@@ -316,10 +323,11 @@ int detect_parker_soc() {
 }
 #endif // __aarch64__
 
-// Should work for CUDA 8.0 through 12.2
+// Should work for CUDA 8.0 through 12.2, plus 12.4 and 12.6
 // 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) {
+        // When the old API is used on GPUs with over 64 TPCs, disable all TPCs >64
         uint128_t full_mask = -1;
         full_mask <<= 64;
         full_mask |= mask;
@@ -377,6 +385,12 @@ void libsmctrl_set_stream_mask_ext(void* stream, uint128_t mask) {
         case 12020:
                 hw_mask_v2 = (void*)(stream_struct_base + CU_12_2_MASK_OFF);
                 break;
+        case 12040:
+                hw_mask_v2 = (void*)(stream_struct_base + CU_12_4_MASK_OFF);
+                break;
+        case 12060:
+                hw_mask_v2 = (void*)(stream_struct_base + CU_12_6_MASK_OFF);
+                break;
 #elif __aarch64__
         case 9000: {
                 // Jetson TX2 offset is slightly different on CUDA 9.0.
@@ -402,6 +416,9 @@ void libsmctrl_set_stream_mask_ext(void* stream, uint128_t mask) {
         case 11040:
                 hw_mask = (struct stream_sm_mask*)(stream_struct_base + CU_11_4_MASK_OFF_JETSON);
                 break;
+        case 12060:
+                hw_mask = (struct stream_sm_mask*)(stream_struct_base + CU_12_6_MASK_OFF_JETSON);
+                break;
 #endif
         }
 
diff --git a/libsmctrl.h b/libsmctrl.h
index a8207b4..eca1f70 100644
--- a/libsmctrl.h
+++ b/libsmctrl.h
@@ -1,5 +1,5 @@
 /**
- * Copyright 2023 Joshua Bakita
+ * Copyright 2024 Joshua Bakita
  * Library to control TPC masks on CUDA launches. Co-opts preexisting debug
  * logic in the CUDA driver library, and thus requires a build with -lcuda.
  */
@@ -21,7 +21,7 @@ extern void libsmctrl_set_global_mask(uint64_t mask);
 // (overrides global mask)
 // @param stream A cudaStream_t (aka CUstream_st*) to apply the mask on
 // @param mask   A bitmask of enabled/disabled TPCs (see Notes on Bitmasks)
-// Supported: CUDA 8.0 - CUDA 12.1
+// Supported: CUDA 8.0 - CUDA 12.2, plus 12.4 and 12.6
 extern void libsmctrl_set_stream_mask(void* stream, uint64_t mask);
 extern void libsmctrl_set_stream_mask_ext(void* stream, uint128_t mask);
 // Set TPC mask for the next kernel launch from the caller's CPU thread
@@ -47,6 +47,10 @@ extern void libsmctrl_set_next_mask(uint64_t mask);
  *
  * Note that the bitwise inversion operator (~, as used above) is very useful,
  * just be sure to apply it to 64-bit integer literals only! (~0x1 != ~0x1ull)
+ *
+ * On GPUs with over 64 TPCs, use the _mask_ext() functions to support 128-bit
+ * masks. If using a 64-bit mask on a GPU with more than 64 TPCs, all TPCs with
+ * IDs over 64 will be disabled.
  */
 
 /* INFORMATIONAL FUNCTIONS */