Add a README and tests for stream masking and next masking

Also rewrite the global masking test to be much more thorough.

10 files changed, 317 insertions, 108 deletions

diff --git a/.gitignore b/.gitignore
index 437f923..553e0fc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,9 @@ libsmctrl.a
 libsmctrl.o
 libsmctrl.so
 libsmctrl_test_gpc_info
+libsmctrl_test_global_mask
+libsmctrl_test_stream_mask
+libsmctrl_test_next_mask
 *.pyc
+*.o
+.gdb_history
diff --git a/Makefile b/Makefile
index 6abc124..b91f6b9 100644
--- a/Makefile
+++ b/Makefile
@@ -14,14 +14,24 @@ libsmctrl.a: libsmctrl.c libsmctrl.h
         ar rcs $@ libsmctrl.o
 
 # Use static linking with tests to avoid LD_LIBRARY_PATH issues
-libsmctrl_test_gpc_info: libsmctrl_test_gpc_info.c libsmctrl.a
+libsmctrl_test_gpc_info: libsmctrl_test_gpc_info.c libsmctrl.a testbench.h
         $(CC) $< -o $@ -g -L. -l:libsmctrl.a $(LDFLAGS)
 
-libsmctrl_test_global_mask: libsmctrl_test_global_mask.cu libsmctrl.a
-        $(NVCC) $< -o $@ -g -L. -l:libsmctrl.a $(LDFLAGS)
+libsmctrl_test_mask_shared.o: libsmctrl_test_mask_shared.cu testbench.h
+        $(NVCC) $< -c -g
 
-tests: libsmctrl_test_gpc_info libsmctrl_test_global_mask
+libsmctrl_test_global_mask: libsmctrl_test_global_mask.c libsmctrl.a libsmctrl_test_mask_shared.o
+        $(NVCC) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS)
+
+libsmctrl_test_stream_mask: libsmctrl_test_stream_mask.c libsmctrl.a libsmctrl_test_mask_shared.o
+        $(NVCC) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS)
+
+libsmctrl_test_next_mask: libsmctrl_test_next_mask.c libsmctrl.a libsmctrl_test_mask_shared.o
+        $(NVCC) $@.c -o $@ libsmctrl_test_mask_shared.o -g -L. -l:libsmctrl.a $(LDFLAGS)
+
+tests: libsmctrl_test_gpc_info libsmctrl_test_global_mask libsmctrl_test_stream_mask libsmctrl_test_next_mask
 
 clean:
         rm -f libsmctrl.so libsmctrl.a libsmctrl_test_gpu_info \
-              libsmctrl_test_global_mask
+              libsmctrl_test_mask_shared.o libmsctrl_test_global_mask \
+              libsmctrl_test_stream_mask libmsctrl_test_next_mask
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..705f2b6
--- /dev/null
+++ b/README.md
@@ -0,0 +1,80 @@
+# libsmctrl: Quick & Easy Hardware Compute Partitioning on NVIDIA GPUs
+
+This library was developed as part of the following paper:
+
+_J. Bakita and J. H. Anderson, "Hardware Compute Partitioning on NVIDIA GPUs", Proceedings of the 29th IEEE Real-Time and Embedded Technology and Applications Symposium, pp. 54-66, May 2023._
+
+Please cite this paper in any work which leverages our library. Here's the BibTeX entry:
+```
+@inproceedings{bakita2023hardware,
+  title={Hardware Compute Partitioning on {NVIDIA} {GPUs}},
+  author={Bakita, Joshua and Anderson, James H},
+  booktitle={Proceedings of the 29th IEEE Real-Time and Embedded Technology and Applications Symposium},
+  year={2023},
+  month={May},
+  pages={54--66},
+  _series={RTAS}
+}
+```
+
+Please see [the paper](https://www.cs.unc.edu/~jbakita/rtas23.pdf) and libsmctrl.h for details and examples of how to use this library.
+We strongly encourage consulting those resources first; the below comments serve merely as an appendum.
+
+## Run Tests
+To test partitioning:
+```
+make tests
+./libsmctrl_test_global_mask
+./libsmctrl_test_stream_mask
+./libsmctrl_test_next_mask
+```
+
+And if `nvdebug` has been installed:
+```
+./libsmctrl_test_gpu_info
+```
+
+## Supported GPUs
+
+#### Known Working
+
+- NVIDIA GPUs from compute capability 3.5 through 8.9, including embedded "Jetson" GPUs
+- CUDA 8.1 through 12.2
+- `x86_64` and Jetson `aarch64` platforms
+
+#### Known Issues
+
+- `next_mask` will not override `stream_mask` on CUDA 12.0+
+- `global_mask` and `next_mask` cannot disable TPCs with IDs above 128
+    - Only relevant on GPUs with over 128 TPCs, such as the RTX 6000 Ada
+- Untested on H100 (compute capability 9.0)
+- Untested on non-Jetson `aarch64` platforms
+
+## Important Limitations
+
+1. Only supports partitioning _within_ a single GPU context.
+   At time of writing, it is challenging to impossible to share a GPU context across multiple CPU address spaces.
+   The implication is that your applications must first be combined together into a single CPU process.
+2. No aspect of this system prevents implicit synchronization on the GPU.
+   See prior work, particularly that of Amert et al. (perhaps the CUPiD^RT paper), for ways to avoid this.
+
+## Porting to New Architectures
+
+Build the tests with `make tests`. And then run the following:
+```
+for (( i=0; $?!=0; i+=8 )); do MASK_OFF=$i ./libsmctrl_test_stream_mask; done
+```
+
+How this works:
+
+1. If `MASK_OFF` is set, `libsmctrl` applies this as a byte offset to a base address for the location
+   of the SM mask fields in CUDA's stream data structure.
+  - That base address is the one for CUDA 12.2 at time of writing
+2. The stream masking test is run.
+3. If the test succeeded (returned zero) the loop aborts, otherwise it increments the offset to attempt and repeats.
+
+Once this loop aborts, take the found offset and add it into the switch statement for the appropriate CUDA version and CPU architecture.
+
+## TODO
+
+- Add a test to check that more-granularly-set compute masks override more-corsely-set ones.
diff --git a/libsmctrl_test_global_mask.c b/libsmctrl_test_global_mask.c
new file mode 100644
index 0000000..9e3189b
--- /dev/null
+++ b/libsmctrl_test_global_mask.c
@@ -0,0 +1,10 @@
+// Copyright 2023 Joshua Bakita
+// Test libsmctrl_set_global_mask().
+// All types of partitioning use the same test, so this file is trival.
+
+#include "libsmctrl_test_mask_shared.h"
+
+int main() {
+  return test_constrained_size_and_location(PARTITION_GLOBAL);
+}
+
diff --git a/libsmctrl_test_global_mask.cu b/libsmctrl_test_global_mask.cu
deleted file mode 100644
index f6a487f..0000000
--- a/libsmctrl_test_global_mask.cu
+++ /dev/null
@@ -1,103 +0,0 @@
-#include <error.h>
-#include <errno.h>
-#include <stdio.h>
-#include <stdbool.h>
-#include <cuda_runtime.h>
-
-#include "libsmctrl.h"
-#include "testbench.h"
-
-__global__ void read_and_store_smid(uint8_t* smid_arr) {
-  if (threadIdx.x != 1)
-    return;
-  int smid;
-  asm("mov.u32 %0, %%smid;" : "=r"(smid));
-  smid_arr[blockIdx.x] = smid;
-}
-
-// Assuming SMs continue to support a maximum of 2048 resident threads, six
-// blocks of 1024 threads should span at least three SMs without partitioning
-#define NUM_BLOCKS 6
-
-int sort_asc(const void* a, const void* b) {
-  return *(uint8_t*)a - *(uint8_t*)b;
-}
-
-// Warning: Mutates input array via qsort
-int count_unique(uint8_t* arr, int len) {
-  qsort(arr, len, 1, sort_asc);
-  int num_uniq = 1;
-  for (int i = 0; i < len - 1; i++)
-    num_uniq += (arr[i] != arr[i + 1]);
-  return num_uniq;
-}
-
-int main() {
-  cudaError_t err; // Needed by SAFE() macro
-  int res;
-  uint8_t *smids_native_d, *smids_native_h;
-  uint8_t *smids_partitioned_d, *smids_partitioned_h;
-  int uniq_native, uniq_partitioned;
-  uint32_t num_tpcs;
-  int num_sms, sms_per_tpc;
-
-  // Determine number of SMs per TPC
-  SAFE(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0));
-  if (res = libsmctrl_get_tpc_info_cuda(&num_tpcs, 0))
-    error(1, res, "libsmctrl_test_global: Unable to get TPC configuration for test");
-  sms_per_tpc = num_sms/num_tpcs;
-
-  // Test baseline (native) behavior without partitioning
-  SAFE(cudaMalloc(&smids_native_d, NUM_BLOCKS));
-  if (!(smids_native_h = (uint8_t*)malloc(NUM_BLOCKS)))
-    error(1, errno, "libsmctrl_test_global: Unable to allocate memory for test");
-  read_and_store_smid<<<NUM_BLOCKS, 1024>>>(smids_native_d);
-  SAFE(cudaMemcpy(smids_native_h, smids_native_d, NUM_BLOCKS, cudaMemcpyDeviceToHost));
-
-  uniq_native = count_unique(smids_native_h, NUM_BLOCKS);
-  if (uniq_native < sms_per_tpc) {
-    printf("libsmctrl_test_global: ***Test failure.***\n"
-           "libsmctrl_test_global: Reason: In baseline test, %d blocks of 1024 "
-           "threads were launched on the GPU, but only %d SMs were utilized, "
-           "when it was expected that at least %d would be used.\n", NUM_BLOCKS,
-           uniq_native, sms_per_tpc);
-    return 1;
-  }
-
-  // Verify that partitioning changes the SMID distribution
-  libsmctrl_set_global_mask(~0x1); // Enable only one TPC
-  SAFE(cudaMalloc(&smids_partitioned_d, NUM_BLOCKS));
-  if (!(smids_partitioned_h = (uint8_t*)malloc(NUM_BLOCKS)))
-    error(1, errno, "libsmctrl_test_global: Unable to allocate memory for test");
-  read_and_store_smid<<<NUM_BLOCKS, 1024>>>(smids_partitioned_d);
-  SAFE(cudaMemcpy(smids_partitioned_h, smids_partitioned_d, NUM_BLOCKS, cudaMemcpyDeviceToHost));
-
-  // Make sure it only ran on the number of TPCs provided
-  // May run on up to two SMs, as up to two per TPC
-  uniq_partitioned = count_unique(smids_partitioned_h, NUM_BLOCKS);
-  if (uniq_partitioned > sms_per_tpc) {
-    printf("libsmctrl_test_global: ***Test failure.***\n"
-           "libsmctrl_test_global: Reason: With global TPC mask set to "
-           "constrain all kernels to a single TPC, a kernel of %d blocks of "
-           "1024 threads was launched and found to run on %d SMs (at most %d---"
-           "one TPC---expected).\n", NUM_BLOCKS, uniq_partitioned, sms_per_tpc);
-    return 1;
-  }
-
-  // Make sure it ran on the right TPC
-  if (smids_partitioned_h[NUM_BLOCKS - 1] > sms_per_tpc - 1) {
-    printf("libsmctrl_test_global: ***Test failure.***\n"
-           "libsmctrl_test_global: Reason: With global TPC mask set to"
-           "constrain all kernels to the first TPC, a kernel was run and found "
-           "to run on an SM ID as high as %d (max of %d expected).\n",
-           smids_partitioned_h[NUM_BLOCKS - 1], sms_per_tpc - 1);
-    return 1;
-  }
-
-  printf("libsmctrl_test_global: Test passed!\n"
-         "libsmctrl_test_global: Reason: With a global partition enabled which "
-         "contained only TPC ID 0, the test kernel was found to use only %d "
-         "SMs (%d without), and all SMs in-use had IDs below %d (were contained"
-         " in the first TPC).\n", uniq_partitioned, uniq_native, sms_per_tpc);
-  return 0;
-}
diff --git a/libsmctrl_test_mask_shared.cu b/libsmctrl_test_mask_shared.cu
new file mode 100644
index 0000000..c3817fa
--- /dev/null
+++ b/libsmctrl_test_mask_shared.cu
@@ -0,0 +1,140 @@
+// Copyright 2023 Joshua Bakita
+#include <error.h>
+#include <errno.h>
+#include <stdio.h>
+#include <stdbool.h>
+#include <cuda_runtime.h>
+
+#include "libsmctrl.h"
+#include "testbench.h"
+#include "libsmctrl_test_mask_shared.h"
+
+__global__ void read_and_store_smid(uint8_t* smid_arr) {
+  if (threadIdx.x != 1)
+    return;
+  int smid;
+  asm("mov.u32 %0, %%smid;" : "=r"(smid));
+  smid_arr[blockIdx.x] = smid;
+}
+
+// Assuming SMs continue to support a maximum of 2048 resident threads, six
+// blocks of 1024 threads should span at least three SMs without partitioning
+#define NUM_BLOCKS 142 //6
+
+static int sort_asc(const void* a, const void* b) {
+  return *(uint8_t*)a - *(uint8_t*)b;
+}
+
+// Warning: Mutates input array via qsort
+static int count_unique(uint8_t* arr, int len) {
+  qsort(arr, len, 1, sort_asc);
+  int num_uniq = 1;
+  for (int i = 0; i < len - 1; i++)
+    num_uniq += (arr[i] != arr[i + 1]);
+  return num_uniq;
+}
+
+// Test that adding an SM mask:
+// 1. Constrains the number of SMs accessible
+// 2. Constrains an application to the correct subset of SMs
+int test_constrained_size_and_location(enum partitioning_type part_type) {
+  int res;
+  uint8_t *smids_native_d, *smids_native_h;
+  uint8_t *smids_partitioned_d, *smids_partitioned_h;
+  int uniq_native, uniq_partitioned;
+  uint32_t num_tpcs;
+  int num_sms, sms_per_tpc;
+  cudaStream_t stream;
+
+  SAFE(cudaStreamCreate(&stream));
+
+  // Determine number of SMs per TPC
+  SAFE(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0));
+  if (res = libsmctrl_get_tpc_info_cuda(&num_tpcs, 0))
+    error(1, res, "Unable to get TPC configuration for test");
+  sms_per_tpc = num_sms/num_tpcs;
+
+  // Test baseline (native) behavior without partitioning
+  SAFE(cudaMalloc(&smids_native_d, NUM_BLOCKS));
+  if (!(smids_native_h = (uint8_t*)malloc(NUM_BLOCKS)))
+    error(1, errno, "Unable to allocate memory for test");
+  read_and_store_smid<<<NUM_BLOCKS, 1024, 0, stream>>>(smids_native_d);
+  SAFE(cudaMemcpy(smids_native_h, smids_native_d, NUM_BLOCKS, cudaMemcpyDeviceToHost));
+
+  uniq_native = count_unique(smids_native_h, NUM_BLOCKS);
+  if (uniq_native < sms_per_tpc) {
+    printf("%s: ***Test failure.***\n"
+           "%s: Reason: In baseline test, %d blocks of 1024 "
+           "threads were launched on the GPU, but only %d SMs were utilized, "
+           "when it was expected that at least %d would be used.\n", program_invocation_name, program_invocation_name, NUM_BLOCKS,
+           uniq_native, sms_per_tpc);
+    return 1;
+  }
+
+  // Test at 32-TPC boundaries to verify that the mask is applied in the
+  // correct order to each of the QMD/stream struct fields.
+  char* reason[4] = {0};
+  for (int enabled_tpc = 0; enabled_tpc < num_tpcs && enabled_tpc < 128; enabled_tpc += 32) {
+    uint128_t mask = 1;
+    mask <<= enabled_tpc;
+    mask = ~mask;
+
+    // Apply partitioning to enable only the first TPC of each 32-bit block
+    switch (part_type) {
+      case PARTITION_GLOBAL:
+        libsmctrl_set_global_mask(mask);
+        break;
+      case PARTITION_STREAM:
+        libsmctrl_set_stream_mask_ext(stream, mask);
+        break;
+      case PARTITION_NEXT:
+        libsmctrl_set_next_mask(mask);
+        break;
+      default:
+        error(1, 0, "Shared test core called with unrecognized partitioning type.");
+    }
+
+    // Verify that partitioning changes the SMID distribution
+    SAFE(cudaMalloc(&smids_partitioned_d, NUM_BLOCKS));
+    if (!(smids_partitioned_h = (uint8_t*)malloc(NUM_BLOCKS)))
+      error(1, errno, "Unable to allocate memory for test");
+    read_and_store_smid<<<NUM_BLOCKS, 1024, 0, stream>>>(smids_partitioned_d);
+    SAFE(cudaMemcpy(smids_partitioned_h, smids_partitioned_d, NUM_BLOCKS, cudaMemcpyDeviceToHost));
+
+    // Make sure it only ran on the number of TPCs provided
+    // May run on up to two SMs, as up to two per TPC
+    uniq_partitioned = count_unique(smids_partitioned_h, NUM_BLOCKS); // Sorts too
+    if (uniq_partitioned > sms_per_tpc) {
+      printf("%s: ***Test failure.***\n"
+             "%s: Reason: With TPC mask set to "
+             "constrain all kernels to a single TPC, a kernel of %d blocks of "
+             "1024 threads was launched and found to run on %d SMs (at most %d---"
+             "one TPC---expected).\n", program_invocation_name, program_invocation_name, NUM_BLOCKS, uniq_partitioned, sms_per_tpc);
+      return 1;
+    }
+
+    // Make sure it ran on the right TPC
+    if (smids_partitioned_h[NUM_BLOCKS - 1] > (enabled_tpc * sms_per_tpc) + sms_per_tpc - 1 ||
+        smids_partitioned_h[NUM_BLOCKS - 1] < (enabled_tpc * sms_per_tpc)) {
+      printf("%s: ***Test failure.***\n"
+             "%s: Reason: With TPC mask set to"
+             "constrain all kernels to TPC %d, a kernel was run and found "
+             "to run on an SM IDs: as high as %d and as low as %d (range of %d to %d expected).\n",
+             program_invocation_name, program_invocation_name, enabled_tpc, smids_partitioned_h[NUM_BLOCKS - 1], smids_partitioned_h[0], enabled_tpc * sms_per_tpc + sms_per_tpc - 1, enabled_tpc * sms_per_tpc);
+      return 1;
+    }
+
+    // Div by 32 via a shift
+    asprintf(&reason[enabled_tpc >> 5],
+         "With a partition enabled which "
+         "contained only TPC ID %d, the test kernel was found to use only %d "
+         "SMs (%d without), and all SMs in-use had IDs between %d and %d (were contained"
+         " in TPC %d).", enabled_tpc, uniq_partitioned, uniq_native, smids_partitioned_h[0], smids_partitioned_h[NUM_BLOCKS - 1], enabled_tpc);
+  }
+
+  printf("%s: Test passed!\n", program_invocation_name);
+  for (int i = 0; i < 4 && reason[i]; i++)
+    printf("%s: Reason %d: %s\n", program_invocation_name, i + 1, reason[i]);
+  return 0;
+}
+
diff --git a/libsmctrl_test_mask_shared.h b/libsmctrl_test_mask_shared.h
new file mode 100644
index 0000000..f95757d
--- /dev/null
+++ b/libsmctrl_test_mask_shared.h
@@ -0,0 +1,16 @@
+// Copyright 2023 Joshua Bakita
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+enum partitioning_type {
+  PARTITION_GLOBAL,
+  PARTITION_STREAM,
+  PARTITION_NEXT,
+};
+
+extern int test_constrained_size_and_location(enum partitioning_type part_type);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/libsmctrl_test_next_mask.c b/libsmctrl_test_next_mask.c
new file mode 100644
index 0000000..7faa9e5
--- /dev/null
+++ b/libsmctrl_test_next_mask.c
@@ -0,0 +1,10 @@
+// Copyright 2023 Joshua Bakita
+// Test libsmctrl_set_global_mask().
+// All types of partitioning use the same test, so this file is trival.
+
+#include "libsmctrl_test_mask_shared.h"
+
+int main() {
+  return test_constrained_size_and_location(PARTITION_NEXT);
+}
+
diff --git a/libsmctrl_test_stream_mask.c b/libsmctrl_test_stream_mask.c
new file mode 100644
index 0000000..063f934
--- /dev/null
+++ b/libsmctrl_test_stream_mask.c
@@ -0,0 +1,10 @@
+// Copyright 2023 Joshua Bakita
+// Test libsmctrl_set_global_mask().
+// All types of partitioning use the same test, so this file is trival.
+
+#include "libsmctrl_test_mask_shared.h"
+
+int main() {
+  return test_constrained_size_and_location(PARTITION_STREAM);
+}
+
diff --git a/testbench.h b/testbench.h
new file mode 100644
index 0000000..dff1211
--- /dev/null
+++ b/testbench.h
@@ -0,0 +1,31 @@
+/* Copyright 2021-2023 Joshua Bakita
+ * Header for miscellaneous experimental helper functions.
+ */
+
+// cudaError_t and CUResult can both safely be cast to an unsigned int
+static __thread unsigned int __SAFE_err;
+
+// The very strange cast in these macros is to satisfy two goals at tension:
+// 1. This file should be able to be included in non-CUDA-using files, and thus
+//    should use no CUDA types outside of this macro.
+// 2. We want to typecheck uses of these macros. The driver and runtime APIs
+//    do not have identical error numbers and/or meanings, so runtime library
+//    calls should use SAFE, and driver library calls should use SAFE_D.
+// These macros allow typechecking, but keep a generic global error variable.
+#define SAFE(x) \
+        if ((*(cudaError_t*)(&__SAFE_err) = (x)) != 0) { \
+                printf("(%s:%d) CUDA error %d: %s i.e. \"%s\" returned by %s. Aborting...\n", \
+                       __FILE__, __LINE__, __SAFE_err, cudaGetErrorName((cudaError_t)__SAFE_err), cudaGetErrorString((cudaError_t)__SAFE_err), #x); \
+                exit(1); \
+        }
+
+#define SAFE_D(x) \
+        if ((*(CUresult*)&(__SAFE_err) = (x)) != 0) { \
+                const char* name; \
+                const char* desc; \
+                cuGetErrorName((CUresult)__SAFE_err, &name); \
+                cuGetErrorString((CUresult)__SAFE_err, &desc); \
+                printf("(%s:%d) CUDA error %d: %s i.e. \"%s\" returned by %s. Aborting...\n", \
+                       __FILE__, __LINE__, __SAFE_err, name, desc, #x); \
+                exit(1); \
+        }