9 files changed, 48 insertions, 11 deletions

diff --git a/.gitignore b/.gitignore
index 553e0fc..7c52826 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,7 +4,9 @@ libsmctrl.so
 libsmctrl_test_gpc_info
 libsmctrl_test_global_mask
 libsmctrl_test_stream_mask
+libsmctrl_test_stream_mask_override
 libsmctrl_test_next_mask
+libsmctrl_test_next_mask_override
 *.pyc
 *.o
 .gdb_history
diff --git a/Makefile b/Makefile
index b91f6b9..dcf44e4 100644
--- a/Makefile
+++ b/Makefile
@@ -26,12 +26,19 @@ libsmctrl_test_global_mask: libsmctrl_test_global_mask.c libsmctrl.a libsmctrl_t
 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_stream_mask_override: libsmctrl_test_stream_mask_override.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
+libsmctrl_test_next_mask_override: libsmctrl_test_next_mask_override.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_stream_mask_override libsmctrl_test_next_mask libsmctrl_test_next_mask_override
 
 clean:
         rm -f libsmctrl.so libsmctrl.a libsmctrl_test_gpu_info \
               libsmctrl_test_mask_shared.o libmsctrl_test_global_mask \
-              libsmctrl_test_stream_mask libmsctrl_test_next_mask
+              libsmctrl_test_stream_mask libmsctrl_test_stream_mask_override \
+              libsmctrl_test_next_mask libmsctrl_test_next_mask_override
diff --git a/README.md b/README.md
index 705f2b6..6e7c75d 100644
--- a/README.md
+++ b/README.md
@@ -44,7 +44,8 @@ And if `nvdebug` has been installed:
 
 #### Known Issues
 
-- `next_mask` will not override `stream_mask` on CUDA 12.0+
+- `next_mask` will not override `stream_mask` on CUDA 11.0+
+    - _As of Feb 2024, a fix for this is coming soon..._
 - `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)
@@ -74,7 +75,3 @@ How this works:
 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_mask_shared.cu b/libsmctrl_test_mask_shared.cu
index c3817fa..f3e8933 100644
--- a/libsmctrl_test_mask_shared.cu
+++ b/libsmctrl_test_mask_shared.cu
@@ -87,9 +87,18 @@ int test_constrained_size_and_location(enum partitioning_type part_type) {
       case PARTITION_STREAM:
         libsmctrl_set_stream_mask_ext(stream, mask);
         break;
+      case PARTITION_STREAM_OVERRIDE:
+        libsmctrl_set_global_mask(~mask);
+        libsmctrl_set_stream_mask_ext(stream, mask);
+        break;
       case PARTITION_NEXT:
         libsmctrl_set_next_mask(mask);
         break;
+      case PARTITION_NEXT_OVERRIDE:
+        libsmctrl_set_global_mask(~mask);
+        libsmctrl_set_stream_mask_ext(stream, ~mask);
+        libsmctrl_set_next_mask(mask);
+        break;
       default:
         error(1, 0, "Shared test core called with unrecognized partitioning type.");
     }
diff --git a/libsmctrl_test_mask_shared.h b/libsmctrl_test_mask_shared.h
index f95757d..5a115b7 100644
--- a/libsmctrl_test_mask_shared.h
+++ b/libsmctrl_test_mask_shared.h
@@ -6,7 +6,9 @@ extern "C" {
 enum partitioning_type {
   PARTITION_GLOBAL,
   PARTITION_STREAM,
+  PARTITION_STREAM_OVERRIDE,
   PARTITION_NEXT,
+  PARTITION_NEXT_OVERRIDE,
 };
 
 extern int test_constrained_size_and_location(enum partitioning_type part_type);
diff --git a/libsmctrl_test_next_mask.c b/libsmctrl_test_next_mask.c
index 7faa9e5..2e4a36a 100644
--- a/libsmctrl_test_next_mask.c
+++ b/libsmctrl_test_next_mask.c
@@ -1,5 +1,5 @@
-// Copyright 2023 Joshua Bakita
-// Test libsmctrl_set_global_mask().
+// Copyright 2024 Joshua Bakita
+// Test libsmctrl_set_next_mask().
 // All types of partitioning use the same test, so this file is trival.
 
 #include "libsmctrl_test_mask_shared.h"
diff --git a/libsmctrl_test_next_mask_override.c b/libsmctrl_test_next_mask_override.c
new file mode 100644
index 0000000..6e2f9d7
--- /dev/null
+++ b/libsmctrl_test_next_mask_override.c
@@ -0,0 +1,10 @@
+// Copyright 2024 Joshua Bakita
+// Test that libsmctrl_set_next_mask() overrides the global and stream masks.
+// 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_OVERRIDE);
+}
+
diff --git a/libsmctrl_test_stream_mask.c b/libsmctrl_test_stream_mask.c
index 063f934..d8895df 100644
--- a/libsmctrl_test_stream_mask.c
+++ b/libsmctrl_test_stream_mask.c
@@ -1,5 +1,5 @@
-// Copyright 2023 Joshua Bakita
-// Test libsmctrl_set_global_mask().
+// Copyright 2024 Joshua Bakita
+// Test libsmctrl_set_stream_mask_ext().
 // All types of partitioning use the same test, so this file is trival.
 
 #include "libsmctrl_test_mask_shared.h"
diff --git a/libsmctrl_test_stream_mask_override.c b/libsmctrl_test_stream_mask_override.c
new file mode 100644
index 0000000..5098d17
--- /dev/null
+++ b/libsmctrl_test_stream_mask_override.c
@@ -0,0 +1,10 @@
+// Copyright 2024 Joshua Bakita
+// Test that libsmctrl_set_stream_mask_ext() overrides the 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_OVERRIDE);
+}
+