threadpool_wrapper.hpp

/*
 * 版权所有 (c) 2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可证 2.0 版（“许可证”）获得许可；
 * 您不得使用此文件，除非符合许可证的规定。
 * 您可以在以下位置获取许可证的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，根据许可证分发的软件是基于“原样”分发的，
 * 不附带任何明示或默示的保证或条件。
 * 请参阅许可证，了解许可证下特定语言的权限和限制。
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <functional>

#include <BS_thread_pool.hpp>

#include <kvikio/nvtx.hpp>

namespace kvikio {

template <typename pool_type>
class thread_pool_wrapper : public pool_type {
 public
 thread_pool_wrapper(unsigned int nthreads) : pool_type{nthreads, worker_thread_init_func} {}

 void reset(unsigned int nthreads) { pool_type::reset(nthreads, worker_thread_init_func); }

 private
 inline static std::function<void()> worker_thread_init_func{[] {
 KVIKIO_NVTX_SCOPED_RANGE("worker thread init", 0, NvtxManager::default_color());
 // 重命名线程池中的工作线程，以提高 nsys-ui 的清晰度。
 // 注意：nsys-ui 目前不支持此 NVTX 功能。
 NvtxManager::rename_current_thread("thread pool");
 }};
};

using BS_thread_pool = thread_pool_wrapper<BS::thread_pool>;

} // namespace kvikio