ONNXRuntime源码阅读(二)线程池
ONNXRuntime的线程池接口在Eigen线程池接口基础之上扩展而来(题外话:TensorFlow中的线程池同样是建立在Eigen线程池基础上),以下是线程池的继承关系,其中 ThreadPoolTempl 是对接口的实现:
在 \(Environment::Initialize()\) 函数中,通过调用 \(onnxruntime::concurrency::CreateThreadPool\) 分别构建算子内和算子间的线程池($ intra_op_thread_pool_ $ & \(inter_op_thread_pool_\))。\(CreateThreadPool\)内部通过调用 \(onnxruntime::concurrency::CreateThreadPoolHelper(onnxruntime::Env *env, OrtThreadPoolParams options)\) ,进而初始化 \(ThreadPool\) 实例,并使用 \(unique\_ptr\) 装饰该实例,作为返回结果。以上只是对类的继承关系进行了简单介绍,下面根据代码中的注释文档解释线程池的内部构造。
ORT线程池在构造上可分为底层和上层,其中底层可以窥探线程内部工作机制(EigenNonBlockingThreadPool.h),上层提供方便易用的接口(threadpool.h)。我们的目的是了解底层原理,因此接下来对这一层进行特别介绍。
ORT线程池派生于Eigen的非阻塞线程池,但是很多细节已经随着迭代而更新很多。ORT主要内容包含以下几点:
- 线程池维护一组系统线程(OS threads),用于执行ThreadPoolTempl::WorkerLoop。每个线程都拥有自己的运行队列(RunQueue),运行队列中都是被push进来的待执行的任务(Task)。主要的工作任务是从队列中弹出一个任务并执行至结束。如果线程的运行队列为空,则线程陷入自旋(spin)等待任务到达,并且尝试从其它线程的运行队列中“偷取”任务来执行,如果没有偷来任务,则阻塞在系统中。在创建线程池时会通过配置标志(flag)和常量 spin_count 来实现这种“spin-then-block”操作;
- 虽然所有的任务(Task)都是简单的 void()->void 函数,但是从概念上来看共有三种不同的类型:
- 通过Schedule()方法从外部提交的一次性任务(one-shot task),被用于支持异步工作。这些任务在并行处理器中使用,但是在测试工具(参考threadpool_test.cc中的一些样例)之外没有被广泛使用;
- 运行 parallel loop 的任务。这些任务在 threadpool.cc 中被定义,并通过 RunParallel->SummonWorkers()提交到运行队列中。每个任务将在内部循环,通过 atoic-fetch-and-add 从用户代码中提取迭代,直到循环完成。这种两层方法让我们将超轻量级的per-iteration-batch工作与管理任务对象的成本更高的per-loop工作分开。
- 运行 parallel section 的任务,这是对上面 parallel loop 方法的扩展,这些任务定义在 RunInParallelSection->SummonWorkers。parallel sections的附加层是进一步的摊销成本的方法:创建任务完成的work可以执行一次,然后在一系列循环中利用。
此外,修改后的 Eigen 线程池有几个方面需要强调:
- 运行队列遵从常规设计,在队头/队尾提供push/pop操作,并且针对拥有运行队列的线程优化PopFront的情况。
(未完待续)