MNN在算子计算时实现了多线程进行加速,对应的多线程编译选项:MNN_USE_THREAD_POOL,默认开启,使用 MNN 内部的无锁线程池实现多线程优化,关闭后,视MNN_OPENMP开关选择OpenMP或关闭多线程优化。(注:MNN 的无锁线程池最多允许两个实例同时使用,即最多供两个模型同时推理使用。参考代码 source/backend/cpu/ThreadPool.cpp 中 MNN_THREAD_POOL_MAX_TASKS 宏的定义。
以Matmul算子的计算为例,代码中与多线程相关的两个宏MNN_CONCURRENCY_BEGIN,MNN_CONCURRENCY_END
ErrorCode CPUMatMul::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
// Fill output by zero if one of inputs is empty.
if (inputs.size() == 2 && outputs.size() == 1 &&
(inputs[0]->elementSize() == 0 || inputs[1]->elementSize() == 0)) {
::memset(outputs[0]->host<char>(), 0, outputs[0]->size());
return NO_ERROR;
}
auto APtr = inputs[0]->host<float>();
auto BPtr = inputs[1]->host<float>();
auto CPtr = outputs[0]->host<float>();
for (auto& f : mPreFunctions) {
// 使用多线程进行前处理加速
MNN_CONCURRENCY_BEGIN(tId, f.second) {
f.first(tId, APtr, BPtr); // f.first是前处理函数句柄,传入参数,调用句柄实现
}
MNN_CONCURRENCY_END();
}
/* 上面宏按照线程池的实现展开后的结果
std::pair<std::function<void(int)>, int> task;
task.second = f.second;
task.first = [&](int tId) {f.first(tId, APtr, BPtr);} ; // 又定义一函数句柄,嵌套调用
此句柄是多线程调用的执行函数, tId是线程的索引
auto cpuBn = (CPUBackend*)backend();
MNN::ThreadPool::enqueue(std::move(task), cpuBn->taskIndex()); // 调用线程池接口,多线程计算, cpuBn->taskIndex()获取多线程实例的索引号。
*/
// 执行计算
mComputer->onExecute();
for (auto& f : mPostFunctions) {
// 计算结果的处理,多线程处理,MNN在进行矩阵相乘时使用了Strassen算法,相比通用矩阵乘降低时间复杂度,使用该方法需要进行矩阵的转置,所以对结算结果也要进行转置
MNN_CONCURRENCY_BEGIN(tId, f.second) {
f.first(tId, APtr, BPtr, CPtr);
}
MNN_CONCURRENCY_END();
}
return NO_ERROR;
}
看完了多线程的调用点,下面分析具体实现:
1. 线程池的创建
在创建Runtime的时候会根据传入的线程数,构造ThreadPool,初始化
CPURuntime::CPURuntime(const Backend::Info& info) {
mStaticAllocator.reset(new BufferAllocator(BufferAllocator::Allocator::createDefault()));
mThreadNumber = info.numThread; // 传入的线程数,算子计算要开启的总线程数
mThreadNumber = std::max(1, mThreadNumber);
mThreadNumber = std::min(mThreadNumber, MAX_THREAD_NUMBER);
mPower = BackendConfig::Power_Normal;
mMemory = BackendConfig::Memory_Normal;
mPrecision = BackendConfig::Precision_Normal;
mFlags = 0;
mFlops = MNNGetCPUFlops(mThreadNumber);
#if defined(__aarch64__) && ENABLE_ARMV82
mIsSupportDot = gCPUInfo.dot;
mIsSupportFp16arith = gCPUInfo.fp16arith;
#endif
if (info.user != nullptr) {
mPrecision = info.user->precision;
mPower = info.user->power;
mMemory = info.user->memory;
mFlags = info.user->flags;
}
#ifdef _OPENMP
switch (mPower) {
case BackendConfig::Power_Low:
MNNSetCPUThreadsMode(MNN_CPU_MODE_LITTLE);
break;
case BackendConfig::Power_High:
MNNSetCPUThreadsMode(MNN_CPU_MODE_POWER_FRI);
break;
default:
break;
}
#endif
#ifdef MNN_USE_THREAD_POOL //仅分析threadpool是实现原理
mThreadNumber = ThreadPool::init(mThreadNumber); // 构造线程池,并且初始化
if (mThreadNumber > 1) {
mTaskIndex = ThreadPool::acquireWorkIndex(); //获取线程池实例的索引
} else {
mTaskIndex = -1;
}
if (mTaskIndex >= 0 && mPower == BackendConfig::Power_High) { //开启高性能模式则立即唤醒所有线程
ThreadPool::active();
}
#endif
}
ThreadPool::ThreadPool(int numberThread) {
mNumberThread = numberThread;
mActiveCount = 0;
mTaskAvailable.resize(MNN_THREAD_POOL_MAX_TASKS);
mTasks.resize(MNN_THREAD_POOL_MAX_TASKS); //线程池最多两个实例
for (int t = 0; t < mTasks.size(); ++t) {
mTaskAvailable[t] = true;
for (int i = 0; i < mNumberThread; ++i) {
mTasks[t].second.emplace_back(new std::atomic_bool{false});
}
}
#ifdef MNN_THREAD_LOCK_CPU
std::vector<int> sortedCPUIDs = sortCPUIDByMaxFrequency(numberThread);
#endif
for (int i = 1; i < mNumberThread; ++i) { // 由默认线程创建新的线程
int threadIndex = i;
#ifdef MNN_THREAD_LOCK_CPU
mWorkers.emplace_back([this, sortedCPUIDs, threadIndex]() {
#else
mWorkers.emplace_back([this, threadIndex]() { // 创建新线程的函数句柄
#endif
#ifdef MNN_THREAD_LOCK_CPU
int res = setSchedAffinity(sortedCPUIDs); // 对新起的线程进行绑核
#endif
while (!mStop) {
while (mActiveCount > 0) {
for (int i = 0; i < MNN_THREAD_POOL_MAX_TASKS; ++i) {
if (*mTasks[i].second[threadIndex]) { // 若Index的线程还未执行计算
mTasks[i].first.first(threadIndex); // 取出传入的计算任务,执行计算
{ *mTasks[i].second[threadIndex] = false; }
}
}
std::this_thread::yield(); // 让渡出cpu时间,供其他线程使用cpu
}
std::unique_lock<std::mutex> _l(mQueueMutex);
mCondition.wait(_l, [this] { return mStop || mActiveCount > 0; });
// 此处使用条件变量,在线程没有计算任务的时候处于阻塞状态,直到得到唤起的通知
}
});
}
}
2. 向线程池传入计算任务
线程池唤醒,在执行算子前调用
void ThreadPool::active() {
if (nullptr == gInstance) {
return;
}
{
std::lock_guard<std::mutex> _l(gInstance->mQueueMutex);
gInstance->mActiveCount++;
}
gInstance->mCondition.notify_all(); // 唤醒线程池
}
void ThreadPool::enqueueInternal(TASK&& task, int index) {
if (mActiveCount == 0) {
for (int i = 0; i < task.second; ++i) {
task.first(i);
}
return;
}
int workSize = task.second;
if (workSize > mNumberThread) {
mTasks[index].first = std::make_pair(
[workSize, &task, this](int tId) {
for (int v = tId; v < workSize; v += mNumberThread) {
task.first(v);
}
},
mNumberThread);
workSize = mNumberThread;
} else {
mTasks[index].first = std::move(task); // 任务放入线程池
}
{
for (int i = 1; i < workSize; ++i) {
*mTasks[index].second[i] = true; // 开启每个线程的计算,标志位设True
}
}
mTasks[index].first.first(0); //主线程执行一个计算任务
bool complete = true;
do {
std::this_thread::yield();
complete = true;
for (int i = 1; i < workSize; ++i) {
if (*mTasks[index].second[i]) { // 如果还有线程没有计算,让渡出主线程的计算资源
complete = false;
break;
}
}
// FUNC_PRINT(notComplete);
} while (!complete);
}