MNN框架源码解读--多线程计算线程池实现分析

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);
}