Introduce pysmctrl: A python interface to libsmctrl

Initially supports the GPU information functions via:
- pysmctrl.get_gpc_info(dev_id)
- pysmctrl.get_tpc_info(dev_id)
- pysmctrl.get_tpc_info_cuda(cuda_dev_id)
All functions are extensively documented. See pysmctrl/__init__.py
for details.

Device partitioning functions have yet to be mapped into Python, as
these will require more testing.

As part of this:
- libsmctrl_get_*_info() functions have been modified to consistently
  return positive error codes.
- libsmctrl_get_tpc_info() now uses nvdebug-style device numbering and
  uses libsmctrl_get_gpc_info() under the covers. This should be more
  reliable.
- libsmctrl_get_tpc_info_cuda() has been introduced as an improved
  version of the old libsmctrl_get_tpc_info() function. This continues
  to use CUDA-style device numbering, but is now resiliant to CUDA
  failures.
- Various minor style improvements in libsmctrl.c

4 files changed, 154 insertions, 36 deletions

diff --git a/.gitignore b/.gitignore
index dcff266..437f923 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@ libsmctrl.a
 libsmctrl.o
 libsmctrl.so
 libsmctrl_test_gpc_info
+*.pyc
diff --git a/libsmctrl.c b/libsmctrl.c
index 640001a..98be1ef 100644
--- a/libsmctrl.c
+++ b/libsmctrl.c
@@ -226,13 +226,15 @@ void libsmctrl_set_stream_mask(void* stream, uint64_t mask) {
                 hw_mask = (struct stream_sm_mask*)(stream_struct_base + CU_12_0_MASK_OFF);
                 break;
         default: {
-                // For experimenting to determine the right mask offset, set MASK_OFF (positive and negative numbers supported)
+                // For experimenting to determine the right mask offset, set the MASK_OFF
+                // environment variable (positive and negative numbers are supported)
                 char* mask_off_str = getenv("MASK_OFF");
                 fprintf(stderr, "libsmctrl: Stream masking unsupported on this CUDA version (%d)!\n", ver);
                 if (mask_off_str) {
                         int off = atoi(mask_off_str);
-                        fprintf(stderr, "libsmctrl: Attempting offset %d on CUDA 11.8 base %#x (total off: %#x)\n", off, CU_11_8_MASK_OFF, CU_11_8_MASK_OFF+off);
-                        hw_mask = (struct stream_sm_mask*)(stream_struct_base + CU_11_8_MASK_OFF + off);
+                        fprintf(stderr, "libsmctrl: Attempting offset %d on CUDA 12.1 base %#x "
+                                        "(total off: %#x)\n", off, CU_12_0_MASK_OFF, CU_12_0_MASK_OFF+off);
+                        hw_mask = (struct stream_sm_mask*)(stream_struct_base + CU_12_0_MASK_OFF + off);
                 } else {
                         return;
                 }}
@@ -242,44 +244,21 @@ void libsmctrl_set_stream_mask(void* stream, uint64_t mask) {
         hw_mask->lower = mask;
 }
 
-int libsmctrl_get_tpc_info(uint32_t* num_tpcs, int dev) {
-        int num_sms;
-        int major;
-        int minor;
-        // TODO: Use nvdebug instead of this hardcoded hack
-        cuDeviceGetAttribute(&num_sms, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev);
-        cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev);
-        cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev);
-        // SM masking only works on sm_35+
-        if (major < 3 || (major == 3 && minor < 5))
-                return -ENOTSUP;
-        // Everything newer than Pascal (as of Hopper) has 2 SMs per TPC, as well
-        // as the P100, which is uniquely sm_60
-        int sms_per_tpc;
-        if (major > 6 || (major == 6 && minor == 0))
-                sms_per_tpc = 2;
-        else
-                sms_per_tpc = 1;
-        // It looks like there may be some upcoming weirdness (TPCs with only one SM?)
-        // with Hopper
-        if (major >= 9)
-                fprintf(stderr, "libsmctrl: WARNING, SM masking is untested on Hopper, and will likely yield incorrect results! Proceed with caution.\n");
-        *num_tpcs = num_sms/sms_per_tpc;
-        return 0;
-}
+/* INFORMATIONAL FUNCTIONS */
 
 // Read an integer from a file in `/proc`
 static int read_int_procfile(char* filename, uint64_t* out) {
         char f_data[18] = {0};
         int fd = open(filename, O_RDONLY);
         if (fd == -1)
-                return -errno;
+                return errno;
         read(fd, f_data, 18);
         close(fd);
         *out = strtoll(f_data, NULL, 16);
         return 0;
 }
 
+// We support up to 12 GPCs per GPU, and up to 16 GPUs.
 static uint64_t tpc_mask_per_gpc_per_dev[16][12];
 // Output mask is vtpc-indexed (virtual TPC)
 int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc, int dev) {
@@ -291,13 +270,14 @@ int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc,
         // Maximum number of GPCs supported for this chip
         snprintf(filename, 100, "/proc/gpu%d/num_gpcs", dev);
         if (err = read_int_procfile(filename, &max_gpcs)) {
-                fprintf(stderr, "libsmctrl: nvdebug module must be loaded into kernel before using libsmctrl_get_gpc_info()\n");
+                fprintf(stderr, "libsmctrl: nvdebug module must be loaded into kernel before "
+                                "using libsmctrl_get_*_info() functions\n");
                 return err;
         }
         // TODO: handle arbitrary-size GPUs
         if (dev > 16 || max_gpcs > 12) {
                 fprintf(stderr, "libsmctrl: GPU possibly too large for preallocated map!\n");
-                return -ERANGE;
+                return ERANGE;
         }
         // Set bit = disabled GPC
         snprintf(filename, 100, "/proc/gpu%d/gpc_mask", dev);
@@ -331,3 +311,52 @@ int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc,
         return 0;
 }
 
+int libsmctrl_get_tpc_info(uint32_t* num_tpcs, int dev) {
+        uint32_t num_gpcs;
+        uint64_t* tpcs_per_gpc;
+        int res;
+        if (res = libsmctrl_get_gpc_info(&num_gpcs, &tpcs_per_gpc, dev))
+                return res;
+        *num_tpcs = 0;
+        for (int gpc = 0; gpc < num_gpcs; gpc++) {
+                *num_tpcs += __builtin_popcountl(tpcs_per_gpc[gpc]);
+        }
+        return 0;
+}
+
+// @param dev Device index as understood by CUDA **can differ from nvdebug idx**
+// This implementation is fragile, and could be incorrect for odd GPUs
+int libsmctrl_get_tpc_info_cuda(uint32_t* num_tpcs, int cuda_dev) {
+        int num_sms, major, minor, res = 0;
+        const char* err_str;
+        if (res = cuInit(0))
+                goto abort_cuda;
+        if (res = cuDeviceGetAttribute(&num_sms, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, cuda_dev))
+                goto abort_cuda;
+        if (res = cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, cuda_dev))
+                goto abort_cuda;
+        if (res = cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, cuda_dev))
+                goto abort_cuda;
+        // SM masking only works on sm_35+
+        if (major < 3 || (major == 3 && minor < 5))
+                return ENOTSUP;
+        // Everything newer than Pascal (as of Hopper) has 2 SMs per TPC, as well
+        // as the P100, which is uniquely sm_60
+        int sms_per_tpc;
+        if (major > 6 || (major == 6 && minor == 0))
+                sms_per_tpc = 2;
+        else
+                sms_per_tpc = 1;
+        // It looks like there may be some upcoming weirdness (TPCs with only one SM?)
+        // with Hopper
+        if (major >= 9)
+                fprintf(stderr, "libsmctrl: WARNING, TPC masking is untested on Hopper,"
+                                " and will likely yield incorrect results! Proceed with caution.\n");
+        *num_tpcs = num_sms/sms_per_tpc;
+        return 0;
+abort_cuda:
+        cuGetErrorName(res, &err_str);
+        fprintf(stderr, "libsmctrl: CUDA call failed due to %s. Failing with EIO...\n", err_str);
+        return EIO;
+}
+
diff --git a/libsmctrl.h b/libsmctrl.h
index f342afa..64ae7a7 100644
--- a/libsmctrl.h
+++ b/libsmctrl.h
@@ -50,15 +50,21 @@ extern void set_sm_mask(uint64_t mask) __attribute__((deprecated("Use libsmctrl_
 
 /* INFORMATIONAL FUNCTIONS */
 
-// Get total number of TPCs on device number `dev`.
-extern int libsmctrl_get_tpc_info(uint32_t* num_tpcs, int dev);
 // Get number of GPCs for devices number `dev`, and a GPC-indexed array
 // containing masks of which TPCs are associated with each GPC.
 // Note that the `nvdebug` module must be loaded to use this function.
-// @param  num_enabled_gpcs Location to store number of GPCs in
-// @param  tpcs_for_gpc     Pointer to store pointer to output buffer at
-// @return 0 on success, error code on error
+// @param  num_enabled_gpcs (out) Location to store number of GPCs in
+// @param  tpcs_for_gpc     (out) Pointer to store pointer to output buffer at
+// @param  dev               (in) `nvdebug` device ID
+// @return 0 on success, `errno`-compatible error code on failure
 extern int libsmctrl_get_gpc_info(uint32_t* num_enabled_gpcs, uint64_t** tpcs_for_gpc, int dev);
+// Get total number of TPCs on device number `dev`. Requires `nvdebug`.
+// @param  num_tpcs        (out) Location to store number of TPCs at
+// @param  dev              (in) `nvdebug` device ID
+// @return 0 on success, `errno`-compatible error code on failure
+extern int libsmctrl_get_tpc_info(uint32_t* num_tpcs, int dev);
+// Identical to above, but for a CUDA device ID. Does not require `nvdebug`.
+extern int libsmctrl_get_tpc_info_cuda(uint32_t* num_tpcs, int cuda_dev);
 
 #ifdef __cplusplus
 }
diff --git a/pysmctrl/__init__.py b/pysmctrl/__init__.py
new file mode 100644
index 0000000..5dc6175
--- /dev/null
+++ b/pysmctrl/__init__.py
@@ -0,0 +1,82 @@
+import ctypes, ctypes.util
+import os
+
+# If this is failing, make sure that the directory containing libsmctrl.so is
+# in your LD_LIBRARY_PATH environment variable. You likely need something like:
+# LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/playpen/jbakita/gpu_subdiv/libsmctrl/
+libsmctrl_path = ctypes.util.find_library("libsmctrl")
+if not libsmctrl_path:
+    libsmctrl_path = __path__[0] + "/../libsmctrl.so"
+libsmctrl = ctypes.CDLL(libsmctrl_path)
+
+def get_gpc_info(device_num):
+    """
+    Obtain list of thread processing clusters (TPCs) enabled for each general
+    processing cluster (GPC) in the specified GPU.
+
+    Parameters
+    ----------
+    device_num : int
+        Which device to obtain information for (starts as 0, order is defined
+        by nvdebug module). May not match CUDA device numbering.
+
+    Returns
+    -------
+    list of int64
+        A list as long as the number of GPCs enabled, where each list entry is
+        a bitmask. A bit set at index `i` indicates that TPC `i` is part of the
+        GPC at that list index. Obtained via GPU register reads in `nvdebug`.
+    """
+    num_gpcs = ctypes.c_uint()
+    tpc_masks = ctypes.pointer(ctypes.c_ulonglong())
+    res = libsmctrl.libsmctrl_get_gpc_info(ctypes.byref(num_gpcs), ctypes.byref(tpc_masks), device_num)
+    if res != 0:
+        print("pysmctrl: Unable to call libsmctrl_get_gpc_info(). Raising error %d..."%res)
+        raise OSError(res, os.strerror(res))
+    return [tpc_masks[i] for i in range(num_gpcs.value)]
+
+def get_tpc_info(device_num):
+    """
+    Obtain a count of the total number of thread processing clusters (TPCs)
+    enabled on the specified GPU.
+
+    Parameters
+    ----------
+    device_num : int
+        Which device to obtain TPC count for (starts as 0, order is defined by
+        `nvdebug` module). May not match CUDA device numbering.
+
+    Returns
+    -------
+    int
+        Count of enabled TPCs. Obtained via GPU register reads in `nvdebug`.
+    """
+    num_tpcs = ctypes.c_uint()
+    res = libsmctrl.libsmctrl_get_tpc_info(ctypes.byref(num_tpcs), device_num)
+    if res != 0:
+        print("pysmctrl: Unable to call libsmctrl_get_tpc_info(). Raising error %d..."%res)
+        raise OSError(res, os.strerror(res))
+    return num_tpcs.value
+
+def get_tpc_info_cuda(device_num):
+    """
+    Obtain a count of the total number of thread processing clusters (TPCs)
+    enabled on the specified GPU.
+
+    Parameters
+    ----------
+    device_num : int
+        Which device to obtain TPC count for, as a CUDA device ID.
+
+    Returns
+    -------
+    int
+        Count of enabled TPCs. Obtained via calculations on data from CUDA.
+    """
+    num_tpcs = ctypes.c_uint()
+    res = libsmctrl.libsmctrl_get_tpc_info_cuda(ctypes.byref(num_tpcs), device_num)
+    if res != 0:
+        print("pysmctrl: Unable to call libsmctrl_get_tpc_info_cuda(). Raising error %d..."%res)
+        raise OSError(res, os.strerror(res))
+    return num_tpcs.value
+