详解MindSpore线程池机制(一)——线程池基本逻辑

100 阅读4分钟

基于任务窃取的分布式调度

下图是Eigen线程池示意图,是一种典型的基于任务窃取的分布式调度:

image.png 主要逻辑:

1.线程池初始化时,一般根据系统的CPU核数创建相应的线程数(或传参数指定线程数)

2.开始一般会有一个不属于线程池内的线程向线程池任务队列提交任务(如随机选择一个队列提交)

3.每个线程在创建时绑定一个Run函数,函数逻辑:先尝试从本地队列(线程自己的任务队列)取任务(一般从头部取,第5条解释),若本地队列没有任务,则尝试从其它线程的队列尾部(第5条解释)偷取任务执行。在满足一定条件时(如长时间未有效获取到任务)进入休眠状态,等待被唤醒

4.线程在执行完某个任务(例如深度学习计算图中某个算子)后,可能产生新的任务,将新任务提交到本地任务队列

5. **之所以每个线程从本地队列取任务是从头部,提交任务也是向头部提交,而偷取别的队列任务是从尾部,是为了提高cache命中率进而提高执行效率**。举例:线程0完成算子A的执行,生成了新任务(算子B),那么很可能算子A的输出数据就是算子B的输入,此时线程0所在的CPU核的cache中还缓存A的输出数据,将算子B的任务提交到本地队列头部,再从头部取任务,大概率会取到算子B任务,则cache中数据可有效利用起来,提高执行效率;别的线程从本线程任务队列尾部偷取任务,避免了刚生成的算子B任务被别的线程抢走执行,导致无法利用cache中数据

6. 任务调度机制设计出发点都是为了效率,上面调度机制好处有几点:

1)分布式调度:没有中央调度器,避免调度器的任务调度能力成为系统瓶颈

2)实现了动态的负载均衡:任务量充足时,所有线程都会忙碌起来,充分利用多核算力

3)避免CPU算力浪费:任务量不足时,获取不到任务的线程会休眠

4)考虑了cache数据亲和:提高cache命中率进而提高任务执行效率

Worker关键代码解析

Worker类就是线程池中的工作线程,只不过对其做了一层封装,下面介绍主要方法(代码位置:mindspore/core/mindrt/src/thread/threadpool.cc):

Worker::CreateThread
void Worker::CreateThread() { thread_ = std::make_unique<std::thread>(&Worker::Run, this); }

该方法用于实例化一个std::thread对象并启动执行,线程生命周期内只执行Worker::Run方法。

Worker::Run
void Worker::Run() {
  if (!core_list_.empty()) {
    SetAffinity();
  }
#if !defined(__APPLE__) && !defined(_MSC_VER)
  static std::atomic_int index = {0};
  (void)pthread_setname_np(pthread_self(), ("KernelThread_" + std::to_string(index++)).c_str());
#endif
#ifdef PLATFORM_86
  // Some CPU kernels need set the flush zero mode to improve performance.
  _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);
  _MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON);
#endif
  while (alive_) {
    if (RunLocalKernelTask()) {
      spin_count_ = 0;
    } else {
      RunOtherKernelTask();
      YieldAndDeactive();
    }
    if (spin_count_ > max_spin_count_) {
      WaitUntilActive();
      spin_count_ = 1;
    }
  }
}

由代码可见,Worker生命周期内执行的是一个循环,流程:

1.调用RunLocalKernelTask尝试从本地任务队列取任务执行,若成功,则将spin_count置0(下文解释)。

2.若没能成功执行本地队列任务(一般是本地队列没任务),则运行RunOtherKernelTask,从其它Worker任务队列中偷取任务执行。然后,调用YieldAndDeactive(内部调用了std::this_thread::yield)让出CPU,给其它线程执行。

3.判断spin_count的值,如果大于max_spin_count_,则调用WaitUntilActive使Worker进入休眠状态,直到被唤醒。

上文中,spin_count是自旋计数的意思。在计算机科学中,自旋(Spinning)指线程在等待某个条件(如锁的释放、新任务的到来)时,不主动让出CPU进入休眠状态,而是通过循环(忙等待)持续检查条件是否满足。自旋的好处是可以避免频繁“休眠-唤醒”带来的性能开销。

由上面代码可见,spin_count记录了线程的自旋次数,当自旋次数超过设定的阈值时就会被休眠。 尽管能够偷取任务来执行,但依然会spin_count++,说明偷取任务这种行为虽然被允许,但偷取的频率需要被控制。