diff --git a/unittest/CMakeLists.txt b/unittest/CMakeLists.txt index da1559d2d..7873a6f0d 100644 --- a/unittest/CMakeLists.txt +++ b/unittest/CMakeLists.txt @@ -39,6 +39,7 @@ endif () set( UTIL_SRCS ${KNOWHERE_SOURCE_DIR}/unittest/utils.cpp + ${KNOWHERE_SOURCE_DIR}/unittest/ThreadChecker.cpp ) set( ALL_INDEXING_SRCS diff --git a/unittest/Helper.h b/unittest/Helper.h index 8e405e3fa..b7eaa8a7a 100644 --- a/unittest/Helper.h +++ b/unittest/Helper.h @@ -35,8 +35,8 @@ constexpr int64_t K = 10; constexpr int64_t PINMEM = 1024 * 1024 * 200; constexpr int64_t TEMPMEM = 1024 * 1024 * 300; constexpr int64_t RESNUM = 2; -constexpr int64_t BUILD_INDEX_OMP_NUM = 3; -constexpr int64_t QUERY_OMP_NUM = 4; +constexpr int64_t BUILD_INDEX_OMP_NUM = 17; +constexpr int64_t QUERY_OMP_NUM = 13; class ParamGenerator { public: diff --git a/unittest/ThreadChecker.cpp b/unittest/ThreadChecker.cpp new file mode 100644 index 000000000..794f75035 --- /dev/null +++ b/unittest/ThreadChecker.cpp @@ -0,0 +1,50 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// 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 "unittest/ThreadChecker.h" + +namespace knowhere::threadchecker { + +#define MAX_THREAD_NUM 100000 +#define MIN_THREAD_NUM 0 +#define THREAD_LENGTH 7 // the length of max_thread_num +1 + +int32_t +GetBuildOmpThread(const Config& conf) { + return CheckKeyInConfig(conf, meta::BUILD_INDEX_OMP_NUM) ? GetMetaBuildIndexOmpNum(conf) : omp_get_max_threads(); +} + +int32_t +GetQueryOmpThread(const Config& conf) { + return CheckKeyInConfig(conf, meta::QUERY_OMP_NUM) ? GetMetaQueryOmpNum(conf) : omp_get_max_threads(); +} + +int32_t +GetThreadNum(int pid) { + std::string cmd = "ps -p " + std::to_string(pid) + " -Tf | wc -l"; + FILE* file = popen(cmd.c_str(), "r"); + char fBuff[THREAD_LENGTH]; + if (fgets(fBuff, sizeof(fBuff), file) == nullptr) { + KNOWHERE_THROW_MSG("could not open the file to get thread number"); + } + pclose(file); + const std::size_t len = strlen(fBuff); + if (len > 0 && fBuff[len - 1] == '\n') { + fBuff[len - 1] = '\0'; + } + int32_t ans = atoi(fBuff); + if (ans < MIN_THREAD_NUM || ans > MAX_THREAD_NUM) { + KNOWHERE_THROW_MSG("thread number is out of control"); + } + return ans; +} + +} // namespace knowhere::threadchecker diff --git a/unittest/ThreadChecker.h b/unittest/ThreadChecker.h new file mode 100644 index 000000000..02341b186 --- /dev/null +++ b/unittest/ThreadChecker.h @@ -0,0 +1,32 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// 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 +#include +#include + +#include "common/Log.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" + +namespace knowhere::threadchecker { + +int32_t +GetBuildOmpThread(const Config& conf); + +int32_t +GetQueryOmpThread(const Config& conf); + +int32_t +GetThreadNum(int pid); + +} // namespace knowhere::threadchecker diff --git a/unittest/test_annoy.cpp b/unittest/test_annoy.cpp index fa1ff8da1..d191e8271 100644 --- a/unittest/test_annoy.cpp +++ b/unittest/test_annoy.cpp @@ -17,6 +17,7 @@ #include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "unittest/Helper.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -146,3 +147,14 @@ TEST_P(AnnoyTest, annoy_slice) { auto result = index_->Query(query_dataset, conf_, nullptr); AssertAnns(result, nq, knowhere::GetMetaTopk(conf_)); } + +TEST_P(AnnoyTest, annoy_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} diff --git a/unittest/test_binaryidmap.cpp b/unittest/test_binaryidmap.cpp index eadbf8dc2..ec575c407 100644 --- a/unittest/test_binaryidmap.cpp +++ b/unittest/test_binaryidmap.cpp @@ -17,6 +17,7 @@ #include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "unittest/Helper.h" #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -268,3 +269,14 @@ TEST_P(BinaryIDMAPTest, binaryidmap_range_search_substructure) { auto qd = knowhere::GenDataset(nq, dim, xq_bin.data()); ASSERT_ANY_THROW(index_->QueryByRange(qd, conf_, nullptr)); } + +TEST_P(BinaryIDMAPTest, binaryidmap_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} diff --git a/unittest/test_binaryivf.cpp b/unittest/test_binaryivf.cpp index 7c8fcd67f..2fc14fa46 100644 --- a/unittest/test_binaryivf.cpp +++ b/unittest/test_binaryivf.cpp @@ -19,6 +19,7 @@ #include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "unittest/Helper.h" #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -228,3 +229,14 @@ TEST_P(BinaryIVFTest, binaryivf_range_search_substructure) { knowhere::SetMetaMetricType(conf_, knowhere::metric::SUBSTRUCTURE); ASSERT_ANY_THROW(index_->Train(base_dataset, conf_)); } + +TEST_P(BinaryIVFTest, binaryivf_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} diff --git a/unittest/test_hnsw.cpp b/unittest/test_hnsw.cpp index 873625401..ebdfd22e2 100644 --- a/unittest/test_hnsw.cpp +++ b/unittest/test_hnsw.cpp @@ -20,6 +20,7 @@ #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "unittest/Helper.h" #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -315,3 +316,14 @@ TEST_P(HNSWTest, HNSW_range_trace_visit) { CheckFederResult(result, nb); } + +TEST_P(HNSWTest, HNSW_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} diff --git a/unittest/test_idmap.cpp b/unittest/test_idmap.cpp index 6416089a3..c97d4d86e 100644 --- a/unittest/test_idmap.cpp +++ b/unittest/test_idmap.cpp @@ -25,6 +25,7 @@ #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" #include "unittest/Helper.h" @@ -258,6 +259,17 @@ TEST_P(IDMAPTest, idmap_range_search_ip) { knowhere::KnowhereConfig::SetBlasThreshold(old_blas_threshold); } +TEST_P(IDMAPTest, idmap_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} + #ifdef KNOWHERE_GPU_VERSION TEST_P(IDMAPTest, idmap_copy) { ASSERT_TRUE(!xb.empty()); diff --git a/unittest/test_ivf.cpp b/unittest/test_ivf.cpp index 7bb084910..9bfa88f8e 100644 --- a/unittest/test_ivf.cpp +++ b/unittest/test_ivf.cpp @@ -34,6 +34,7 @@ #include "unittest/Helper.h" #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -216,6 +217,17 @@ TEST_P(IVFTest, ivf_range_search_ip) { } } +TEST_P(IVFTest, ivf_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +} + // TODO(linxj): deprecated #ifdef KNOWHERE_GPU_VERSION TEST_P(IVFTest, clone_test) { diff --git a/unittest/test_ivf_nm.cpp b/unittest/test_ivf_nm.cpp index 1d8db96e0..45b1b4a58 100644 --- a/unittest/test_ivf_nm.cpp +++ b/unittest/test_ivf_nm.cpp @@ -34,6 +34,7 @@ #include "knowhere/utils/distances_simd.h" #include "unittest/Helper.h" #include "unittest/range_utils.h" +#include "unittest/ThreadChecker.h" #include "unittest/utils.h" using ::testing::Combine; @@ -258,4 +259,16 @@ TEST_P(IVFNMTest, ivfnm_get_meta) { knowhere::Config j2 = nlohmann::json::parse(json_id_set); ASSERT_NO_THROW(nlohmann::from_json(j2, id_set)); std::cout << "id_set num = " << id_set.size() << std::endl; -} \ No newline at end of file +} + +TEST_P(IVFNMTest, ivfnm_omp) { + int pid = getpid(); + int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid); + index_->BuildAll(base_dataset, conf_); + int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1); + LoadRawData(index_, base_dataset, conf_); + auto result = index_->Query(query_dataset, conf_, nullptr); + int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid); + EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1); +}