From 2c4b2d784815c5a2b4c49592b912c043f3d2a954 Mon Sep 17 00:00:00 2001 From: Joshua Bakita Date: Tue, 17 Oct 2023 15:01:34 -0400 Subject: Support global masking on x86_64 and aarch64 with CUDA 10.2 Also improve documentation and abort with an error message if attempting to set a global SM mask on an unsupported CUDA version. (Would crash/corrupt state before.) Also uncomment a line which errantly disabled global masking on CUDA 10.2 on aarch64. Tested with CUDA 10.2 on: - x86_64 (GTX 1060 3GB, driver 440.100, jbakita-old.cs.unc.edu) - aarch64 (Jetson TX2, driver r32.5, grizzly.cs.unc.edu) --- libsmctrl.c | 122 ++++++++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 95 insertions(+), 27 deletions(-) diff --git a/libsmctrl.c b/libsmctrl.c index 94578a1..dfd71b8 100644 --- a/libsmctrl.c +++ b/libsmctrl.c @@ -6,12 +6,18 @@ #include #include +#include #include #include #include #include -// Layout of mask control fields in CUDA's `globals` struct +// In functions that do not return an error code, we favor terminating with an +// error rather than merely printing a warning and continuing. +#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; @@ -19,42 +25,100 @@ struct global_sm_control { /*** CUDA Globals Manipulation. CUDA 10.2 only ***/ -// Ends up being 0x7fb7fa3408 in some binaries +// 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. + * 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_sm_control_10() { +static void setup_g_sm_control_10() { if (g_sm_control) return; - // Defeat relocation. cudbgReportDriverApiErrorFlags is relocated by - // the loader, but not subject to ASLR (it's always at a constant - // offset in the loaded instance of libcuda.so). Our target is also at - // a constant offset, so we can use the address of - // cudbgReportDriverApiErrorFlags as a reference point. - // Note: cudbgReportDriverApiErrorFlags is currently the closest known - // symbol to **the table**. cudbgDebuggerInitialized is the closest to - // globals itself (+7424 == SM mask control), but we perfer the table - // lookup approach for now, as that's what cuDeviceGetCount() does. + // 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.) extern uint32_t cudbgReportDriverApiErrorFlags; - uint32_t* sym = 0;//&cudbgReportDriverApiErrorFlags; - // In some binaries, the following works out to 0x7fb7ea6000, and - // that's what shows up in the adrp instruction in cuDeviceGetCount() - // in the lead-up to get globals.numDevices. Find this offset by - // calling cuDeviceGetCount(0xdeadbeef), catching the segfault in GDB, - // disassembling the prior instructions, taking the adrp constant, and - // subtracting the address of cudbgReportDriverApiErrorFlags from it. + uint32_t* sym = &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?) - uintptr_t globals_addr = *(tbl_base + 459); // Offset 3672 on aarch64 + // 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-NULL SM disable mask during setup! g_sm_control is likely invalid---use at own risk.\n"); + 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+ ***/ @@ -119,19 +183,23 @@ static void setup_sm_control_11() { fprintf(stderr, "libsmctrl: Error enabling launch callback. Error %d\n", res); } -// Common masking control +// Set default mask for all launches void libsmctrl_set_global_mask(uint64_t mask) { int ver; cuDriverGetVersion(&ver); - if (ver <= 10020) { + if (ver == 10020) { if (!g_sm_control) - setup_sm_control_10(); + setup_g_sm_control_10(); g_sm_control->mask = mask; g_sm_control->enabled = 1; - } else { + } else if (ver > 10020) { if (!sm_control_setup_called) setup_sm_control_11(); g_sm_mask = mask; + } else { // < CUDA 10.2 + abort(1, ENOSYS, "Global masking requires at least CUDA 10.2; " + "this application is using CUDA %d.%d", + ver / 1000, (ver % 100)); } } -- cgit v1.2.2