Bugfix stream-mask override, support old CUDA, and start Hopper support

Use a different callback to intercept the TMD/QMD later in the
launch pipeline.

Major improvements:
- Fix bug with next mask not overriding stream mask on CUDA 11.0+
- Add CUDA 6.5-10.2 support for next- and global-granularity
  partitioning masks on x86_64 and aarch64 Jetson
- Remove libdl dependency
- Partially support TMD/QMD Version 4 (Hopper)

Minor improvements:
- Check for sufficient CUDA version before before attempting to
  apply a next-granularity partitioning mask
- Only check for sufficient CUDA version on the first call to
  `libsmctrl_set_next_mask()` or `libsmctrl_set_global_mask()`,
  rather than checking every time (lowers overheads)
- Check that TMD version is sufficient before modifying it
- Improve documentation

Issues:
- Partitioning mask bits have a different meaning in TMD/QMD
  Version 4 and require floorsweeping and remapping information to
  properly construct. This information will be forthcoming in
  future releases of libsmctrl and nvdebug.

1 files changed, 86 insertions, 161 deletions

diff --git a/libsmctrl.c b/libsmctrl.c
index 1018e44..24a3177 100644
--- a/libsmctrl.c
+++ b/libsmctrl.c
@@ -5,31 +5,34 @@
  *
  * This file implements partitioning via three different mechanisms:
  * - Modifying the QMD/TMD immediately prior to upload
- * - Changing a field in CUDA's global struct that CUDA applies to the QMD/TMD
  * - Changing a field in CUDA's stream struct that CUDA applies to the QMD/TMD
  * This table shows the mechanism used with each CUDA version:
  *   +-----------+---------------+---------------+--------------+
  *   |  Version  |  Global Mask  |  Stream Mask  |  Next Mask   |
  *   +-----------+---------------+---------------+--------------+
- *   | 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          |
+ *   | 8.0-12.6  | 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.
  * 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
+ * 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`.
  */
 #include <cuda.h>
 
 #include <errno.h>
 #include <error.h>
 #include <fcntl.h>
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <unistd.h>
 
-#include <dlfcn.h>
-
 #include "libsmctrl.h"
 
 // In functions that do not return an error code, we favor terminating with an
@@ -37,148 +40,66 @@
 #define abort(ret, errno, ...) error_at_line(ret, errno, __FILE__, __LINE__, \
                                              __VA_ARGS__)
 
-// Layout of mask control fields to match CUDA's static global struct
-struct global_sm_control {
-        uint32_t enabled;
-        uint64_t mask;
-} __attribute__((packed));
-
-/*** CUDA Globals Manipulation. CUDA 10.2 only ***/
-
-// Ends up being 0x7fb7fa3408 in some binaries (CUDA 10.2, Jetson)
-static struct global_sm_control* g_sm_control = NULL;
-
-/* Find the location of CUDA's `globals` struct and the SM mask control fields
- * No symbols are exported from within `globals`, so this has to do a very
- * messy lookup, following the pattern of the assembly of `cuDeviceGetCount()`.
- * Don't call this before the CUDA library has been initialized.
- * (Note that this appears to work, even if built on CUDA > 10.2.)
- */
-static void setup_g_sm_control_10() {
-        if (g_sm_control)
-                return;
-        // The location of the static global struct containing the global SM
-        // mask field will vary depending on where the loader locates the CUDA
-        // library. In order to reliably modify this struct, we must defeat
-        // that relocation by deriving its location relative to a known
-        // reference point.
-        //
-        // == Choosing a Reference Point:
-        // The cudbg* symbols appear to be relocated to a constant offset from
-        // the globals structure, and so we use the address of the symbol
-        // `cudbgReportDriverApiErrorFlags` as our reference point. (This ends
-        // up being the closest to an intermediate table we use as part of our
-        // lookup---process discussed below.)
-        //
-        // Unfortunately, the symbol we reference is errantly omitted from the
-        // libcuda.so stub used by nvcc starting around CUDA 11.8, so we have to
-        // use dlsym to avoid build-time issues.
-        void* hndl = dlopen(NULL, RTLD_LAZY);
-        uint32_t* sym = dlsym(hndl, "cudbgReportDriverApiErrorFlags");
-
-        // == Deriving Location:
-        // The number of CUDA devices available is co-located in the same CUDA
-        // globals structure that we aim to modify the SM mask field in. The
-        // value in that field can be assigned to a user-controlled pointer via
-        // the cuDeviceGetCount() CUDA Driver Library function. To determine
-        // the location of thu structure, we pass a bad address to the function
-        // and dissasemble the code adjacent to where it segfaults. On the
-        // Jetson Xavier with CUDA 10.2, the assembly is as follows:
-        //   (reg x19 contains cuDeviceGetCount()'s user-provided pointer)
-        //   ...
-        //   0x0000007fb71454b4:  cbz   x19, 0x7fb71454d0 // Check ptr non-zero
-        //   0x0000007fb71454b8:  adrp  x1, 0x7fb7ea6000 // Addr of lookup tbl
-        //   0x0000007fb71454bc:  ldr   x1, [x1,#3672] // Get addr of globals
-        //   0x0000007fb71454c0:  ldr   w1, [x1,#904] // Get count from globals
-        //   0x0000007fb71454c4:  str   w1, [x19] // Store count at user addr
-        //   ...
-        // In this assembly, we can identify that CUDA uses an internal lookup
-        // table to identify the location of the globals structure (pointer
-        // 459 in the table; offset 3672). After obtaining this pointer, it
-        // advances to offset 904 in the global structure, dereferences the
-        // value stored there, and then attempts to store it at the user-
-        // -provided address (register x19). This final line will trigger a
-        // segfault if a non-zero bad address is passed to cuDeviceGetCount().
-        //
-        // On x86_64:
-        //   (reg %rbx contains cuDeviceGetCount()'s user-provided pointer)
-        //   ...
-        //   0x00007ffff6cac01f:  test  %rbx,%rbx // Check ptr non-zero
-        //   0x00007ffff6cac022:  je    0x7ffff6cac038 // ''
-        //   0x00007ffff6cac024:  mov   0x100451d(%rip),%rdx # 0x7ffff7cb0548 // Get globals base address from offset from instruction pointer
-        //   0x00007ffff6cac02b:  mov   0x308(%rdx),%edx // Take globals base address, add an offset of 776, and dereference
-        //   0x00007ffff6cac031:  mov   %edx,(%rbx) // Store count at user addr
-        //   ...
-        // Note that this does not use an intermediate lookup table.
-        //
-        // [Aside: cudbgReportDriverApiErrorFlags is currently the closest
-        // symbol to **the lookup table**. cudbgDebuggerInitialized is closer
-        // to the globals struct itself (+7424 == SM mask control), but we
-        // perfer the table lookup approach for now, as that's what
-        // cuDeviceGetCount() does.]
-
-#if __aarch64__
-        // In my test binary, the lookup table is at address 0x7fb7ea6000, and
-        // this is 1029868 bytes before the address for
-        // cudbgReportDriverApiErrorFlags. Use this information to derive the
-        // location of the lookup in our binary (defeat relocation).
-        uintptr_t* tbl_base = (uintptr_t*)((uintptr_t)sym - 1029868);
-        // Address of `globals` is at offset 3672 (entry 459?) in the table
-        uintptr_t globals_addr = *(tbl_base + 459);
-        // SM mask control is at offset 4888 in the `globals` struct
-        // [Device count at offset 904 (0x388)]
-        g_sm_control = (struct global_sm_control*)(globals_addr + 4888);
-#endif // __aarch64__
-#if __x86_64__
-        // In my test binary, globals is at 0x7ffff7cb0548, which is 1103576
-        // bytes before the address for cudbgReportDriverApiErrorFlags
-        // (0x7ffff7dbdc20). Use this offset to defeat relocation.
-        uintptr_t globals_addr = *(uintptr_t*)((uintptr_t)sym - 1103576);
-        // SM mask control is at offset 4728 in the `globals` struct
-        // [Device count at offset 776 (0x308)]
-        g_sm_control = (struct global_sm_control*)(globals_addr + 4728);
-#endif // __x86_64__
-        // SM mask should be empty by default
-        if (g_sm_control->enabled || g_sm_control->mask)
-                fprintf(stderr, "Warning: Found non-empty SM disable mask "
-                                "during setup! libsmctrl_set_global_mask() is "
-                                "unlikely to work on this platform!\n");
-}
-
-/*** QMD/TMD-based SM Mask Control via Debug Callback. CUDA 11+ ***/
+/*** QMD/TMD-based SM Mask Control via Debug Callback. ***/
 
-// Tested working on CUDA x86_64 11.0-12.2.
-// Tested not working on aarch64 or x86_64 10.2
+// Tested working on x86_64 CUDA 6.5, 9.1, and various 10+ versions
+// (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};
-#define LAUNCH_DOMAIN 0x3
-#define LAUNCH_PRE_UPLOAD 0x3
+// 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
+// Global mask (applies across all threads)
 static uint64_t g_sm_mask = 0;
+// Next mask (applies per-thread)
 static __thread uint64_t g_next_sm_mask = 0;
-static char sm_control_setup_called = 0;
-static void launchCallback(void *ukwn, int domain, int cbid, const void *in_params) {
-        // The third 8-byte element in `in_parms` is a pointer to the stream struct.
-        // This exists even when in_params < 0x50. This could be used to implement
-        // stream masking without the manual offsets specified elsewhere (store a
-        // table of stream pointers to masks and do a lookup here).
+// Flag value to indicate if setup has been completed
+static bool sm_control_setup_called = false;
+
+// 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.
+static void control_callback_v2(void *ukwn, int domain, int cbid, const void *in_params) {