Include multiple label types

Both in tests and benchmarks
NVIDIA · Oct 18, 2023 · 7f1d083 · 7f1d083
1 parent 6184e50
commit 7f1d083
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/benchmarks/trie/insert_bench.cu b/benchmarks/trie/insert_bench.cu
@@ -29,12 +29,12 @@ using namespace cuco::utility;
 /**
  * @brief A benchmark evaluating `cuco::experimental::trie::insert` performance
  */
-void trie_insert(nvbench::state& state)
+template <typename LabelType>
+void trie_insert(nvbench::state& state, nvbench::type_list<LabelType>)
 {
   auto const num_keys       = state.get_int64_or_default("NumKeys", 100 * 1000);
   auto const max_key_length = state.get_int64_or_default("MaxKeyLength", 10);
 
-  using LabelType = int;
   cuco::experimental::trie<LabelType> trie;
 
   thrust::host_vector<LabelType> labels;
@@ -54,7 +54,7 @@ void trie_insert(nvbench::state& state)
   });
 }
 
-NVBENCH_BENCH(trie_insert)
+NVBENCH_BENCH_TYPES(trie_insert, NVBENCH_TYPE_AXES(nvbench::type_list<char, int>))
   .set_name("trie_insert")
   .set_max_noise(defaults::MAX_NOISE)
   .add_int64_axis("NumKeys", std::vector<nvbench::int64_t>{100 * 1000, 1000 * 1000})

diff --git a/benchmarks/trie/lookup_bench.cu b/benchmarks/trie/lookup_bench.cu
@@ -29,12 +29,12 @@ using namespace cuco::utility;
 /**
  * @brief A benchmark evaluating `cuco::experimental::trie::lookup` performance
  */
-void trie_lookup(nvbench::state& state)
+template <typename LabelType>
+void trie_lookup(nvbench::state& state, nvbench::type_list<LabelType>)
 {
   auto const num_keys       = state.get_int64_or_default("NumKeys", 100 * 1000);
   auto const max_key_length = state.get_int64_or_default("MaxKeyLength", 10);
 
-  using LabelType = int;
   cuco::experimental::trie<LabelType> trie;
 
   thrust::host_vector<LabelType> labels;
@@ -62,9 +62,8 @@ void trie_lookup(nvbench::state& state)
   });
 }
 
-NVBENCH_BENCH(trie_lookup)
+NVBENCH_BENCH_TYPES(trie_lookup, NVBENCH_TYPE_AXES(nvbench::type_list<char, int>))
   .set_name("trie_lookup")
   .set_max_noise(defaults::MAX_NOISE)
-  .add_int64_axis("NumKeys",
-                  std::vector<nvbench::int64_t>{100 * 1000, 1000 * 1000, 10 * 1000 * 1000})
+  .add_int64_axis("NumKeys", std::vector<nvbench::int64_t>{100 * 1000, 1000 * 1000})
   .add_int64_axis("MaxKeyLength", std::vector<nvbench::int64_t>{4, 8, 16});
diff --git a/tests/trie/lookup_test.cu b/tests/trie/lookup_test.cu
@@ -27,10 +27,9 @@
 
 using namespace cuco::utility;
 
-TEST_CASE("Lookup test", "")
+template <typename LabelType>
+void trie_lookup_test()
 {
-  using LabelType = int;
-
   std::size_t num_keys       = 64 * 1024;
   std::size_t max_key_length = 6;
 
@@ -56,3 +55,9 @@ TEST_CASE("Lookup test", "")
   trie.lookup(d_labels.begin(), d_offsets.begin(), d_offsets.end(), result.begin());
   REQUIRE(cuco::test::all_of(result.begin(), result.end(), cuco::test::trie::valid_key(num_keys)));
 }
+
+TEST_CASE("Trie lookup", "")
+{
+  trie_lookup_test<int>();
+  trie_lookup_test<char>();
+}