From 279f038494b74d714572149273546f42b85ae5de Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Mon, 18 Sep 2023 14:47:27 -0700
Subject: [PATCH] Add compute_hash_table_size helper function

---
 include/cuco/utility/helper_functions.hpp | 56 +++++++++++++++++++++++
 tests/CMakeLists.txt                      |  1 +
 tests/utility/helper_test.cu              | 32 +++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 include/cuco/utility/helper_functions.hpp
 create mode 100644 tests/utility/helper_test.cu

diff --git a/include/cuco/utility/helper_functions.hpp b/include/cuco/utility/helper_functions.hpp
new file mode 100644
index 000000000..bbe23bb5c
--- /dev/null
+++ b/include/cuco/utility/helper_functions.hpp
@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuco/detail/error.hpp>
+#include <cuco/extent.cuh>
+
+namespace cuco {
+
+/// The default occupancy percentage: 0.5 implies a 50% occupancy
+constexpr double DEFAULT_HASH_TABLE_OCCUPANCY = 0.5;
+
+/**
+ * @brief Computes requisite size of hash table
+ *
+ * Computes the number of entries required in a hash table to satisfy
+ * inserting a specified number of keys to achieve the specified hash table
+ * occupancy.
+ *
+ * @note The output of this helper function for a static extent is not a build-time constant.
+ *
+ * @throw If the desired occupancy is no bigger than zero
+ * @throw If the desired occupancy is larger than one
+ *
+ * @param size The number of elements that will be inserted
+ * @param desired_occupancy The desired occupancy percentage, e.g., 0.5 implies a
+ * 50% occupancy
+ *
+ * @return The size of the hash table that will satisfy the desired occupancy for the specified
+ * input size
+ */
+constexpr std::size_t compute_hash_table_size(
+  std::size_t size, double desired_occupancy = DEFAULT_HASH_TABLE_OCCUPANCY)
+{
+  CUCO_EXPECTS(desired_occupancy > 0., "Desired occupancy must be larger than zero");
+  CUCO_EXPECTS(desired_occupancy <= 1., "Desired occupancy cannot be larger than one");
+
+  // Calculate size of hash map based on the desired occupancy
+  return static_cast<std::size_t>(static_cast<double>(size) / desired_occupancy);
+}
+
+}  // namespace cuco
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 3deeeddf1..cbd680def 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -51,6 +51,7 @@ endfunction(ConfigureTest)
 # - utility tests ---------------------------------------------------------------------------------
 ConfigureTest(UTILITY_TEST
     utility/extent_test.cu
+    utility/helper_test.cu
     utility/storage_test.cu
     utility/fast_int_test.cu
     utility/hash_test.cu)
diff --git a/tests/utility/helper_test.cu b/tests/utility/helper_test.cu
new file mode 100644
index 000000000..dfa39ae46
--- /dev/null
+++ b/tests/utility/helper_test.cu
@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuco/utility/helper_functions.hpp>
+
+#include <catch2/catch_template_test_macros.hpp>
+
+TEMPLATE_TEST_CASE_SIG(
+  "Helper function tests", "", ((typename SizeType), SizeType), (int32_t), (int64_t), (std::size_t))
+{
+  SizeType constexpr num = 89;
+  auto const occupancy   = 0.89;
+
+  auto const result = cuco::compute_hash_table_size(num, occupancy);
+
+  SizeType constexpr gold_reference = 100;
+
+  REQUIRE(result == gold_reference);
+}