tests: Very basic test parsing to run individual tests in test-memory

This may be too much infrastructure for a single unit test, but it makes it
a lot easier to complex single tests without getting lost in the noice of
the other tests that already work.

Branch: HybridCache

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

Author: gabe-l-hart

tests/test-memory.cpp

@@ -1,3 +1,17 @@
+/*------------------------------------------------------------------------------
+ * Unit tests for llama-memory.h and derived memory implementations. It contains
+ * a number of tests which can be run all together or separately.
+ *
+ * USAGE: ./bin/test-memory <test_name1> <test_name2>
+ *
+ * When adding a new test, do the following:
+ *
+ *   1. Add the new test_<memory_type>_description function under the
+ *      appropriate memory type section
+ *
+ *   2. Add `RUN_TEST(test_<memory_type>_description);` to main
+ *----------------------------------------------------------------------------*/
+
 #include "../src/llama-arch.h"
 #include "../src/llama-hparams.h"
 #include "../src/llama-impl.h"
@@ -61,13 +75,22 @@ struct log_scope {
     }
 };
 
-#define LOG_SCOPE() log_scope __log_scope(__func__)
+#define RUN_TEST(test_name)                                                \
+    do {                                                                   \
+        bool run_test = argc < 2;                                          \
+        std::vector<std::string> args(argv + 1, argv + argc);              \
+        if (std::find(args.begin(), args.end(), #test_name) != args.end()) \
+            run_test = true;                                               \
+        if (run_test) {                                                    \
+            log_scope __log_scope(#test_name);                             \
+            test_name();                                                   \
+        }                                                                  \
+    } while (0)
 
 /*- Unified Cache ------------------------------------------------------------*/
 
 /* Test that the unified cache can be constructed and destructed safely */
 static void test_llama_kv_cache_unified_constructor() {
-    LOG_SCOPE();
     auto model = _make_model();
     llama_kv_cache_unified cache(
         /* model   */ *model,
@@ -84,7 +107,6 @@ static void test_llama_kv_cache_unified_constructor() {
 
 /* Test that the recurrent cache can be constructed and destructed safely */
 static void test_llama_kv_cache_recurrent_constructor() {
-    LOG_SCOPE();
     auto model = _make_model(LLM_ARCH_MAMBA);
     llama_kv_cache_recurrent cache(
         /* model   */ *model,
@@ -99,7 +121,6 @@ static void test_llama_kv_cache_recurrent_constructor() {
 
 /* Test that the hybrid cache can be constructed and destructed safely */
 static void test_llama_kv_cache_hybrid_constructor() {
-    LOG_SCOPE();
     auto model = _make_model(
         /* arch          =*/ LLM_ARCH_LLAMA,
         /* n_layer       =*/ 4,
@@ -159,12 +180,13 @@ static void test_llama_kv_cache_hybrid_constructor() {
 
 /*- Main ---------------------------------------------------------------------*/
 
-int main() {
+int main(int argc, char* argv[]) {
     // Unified Cache Tests
-    test_llama_kv_cache_unified_constructor();
+    RUN_TEST(test_llama_kv_cache_unified_constructor);
+    RUN_TEST(test_llama_kv_cache_unified_single_seq);
     // Recurrent Cache Tests
-    test_llama_kv_cache_recurrent_constructor();
+    RUN_TEST(test_llama_kv_cache_recurrent_constructor);
     // Hybrid Cache Tests
-    test_llama_kv_cache_hybrid_constructor();
+    RUN_TEST(test_llama_kv_cache_hybrid_constructor);
     return 0;
 }