dmlc--dgl
Lingfan Yu
a1d50f0f53
* include/dgl/runtime * include * src/runtime * src/graph * src/scheduler * src * clean up CMakeLists * further clean up in cmake * install commands * python/dgl/_ffi/_cython * python/dgl/_ffi/_ctypes * python/dgl/_ffi * python/dgl * some fix * copy right
86 行
2.3 KiB
C++
86 行
2.3 KiB
C++
/*!
|
|
* Copyright (c) 2018 by Contributors
|
|
* \file dgl/runtime/threading_backend.h
|
|
* \brief Utilities for manipulating thread pool threads.
|
|
*/
|
|
#ifndef DGL_RUNTIME_THREADING_BACKEND_H_
|
|
#define DGL_RUNTIME_THREADING_BACKEND_H_
|
|
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace dgl {
|
|
namespace runtime {
|
|
namespace threading {
|
|
|
|
/*!
|
|
* \brief A platform-agnostic abstraction for managing a collection of
|
|
* thread pool threads.
|
|
*/
|
|
class ThreadGroup {
|
|
public:
|
|
class Impl;
|
|
|
|
/*!
|
|
* \brief Creates a collection of threads which run a provided function.
|
|
*
|
|
* \param num_workers The total number of worker threads in this group.
|
|
Includes main thread if `exclude_worker0 = true`
|
|
* \param worker_callback A callback which is run in its own thread.
|
|
Receives the worker_id as an argument.
|
|
* \param exclude_worker0 Whether to use the main thread as a worker.
|
|
* If `true`, worker0 will not be launched in a new thread and
|
|
* `worker_callback` will only be called for values >= 1. This
|
|
* allows use of the main thread as a worker.
|
|
*/
|
|
ThreadGroup(int num_workers,
|
|
std::function<void(int)> worker_callback,
|
|
bool exclude_worker0 = false);
|
|
~ThreadGroup();
|
|
|
|
/*!
|
|
* \brief Blocks until all non-main threads in the pool finish.
|
|
*/
|
|
void Join();
|
|
|
|
enum AffinityMode : int {
|
|
kBig = 1,
|
|
kLittle = -1,
|
|
};
|
|
|
|
/*!
|
|
* \brief configure the CPU id affinity
|
|
*
|
|
* \param mode The preferred CPU type (1 = big, -1 = little).
|
|
* \param nthreads The number of threads to use (0 = use all).
|
|
* \param exclude_worker0 Whether to use the main thread as a worker.
|
|
* If `true`, worker0 will not be launched in a new thread and
|
|
* `worker_callback` will only be called for values >= 1. This
|
|
* allows use of the main thread as a worker.
|
|
*
|
|
* \return The number of workers to use.
|
|
*/
|
|
int Configure(AffinityMode mode, int nthreads, bool exclude_worker0);
|
|
|
|
private:
|
|
Impl* impl_;
|
|
};
|
|
|
|
/*!
|
|
* \brief Platform-agnostic no-op.
|
|
*/
|
|
void Yield();
|
|
|
|
/*!
|
|
* \return the maximum number of effective workers for this system.
|
|
*/
|
|
int MaxConcurrency();
|
|
|
|
|
|
} // namespace threading
|
|
} // namespace runtime
|
|
} // namespace dgl
|
|
|
|
#endif // DGL_RUNTIME_THREADING_BACKEND_H_
|