下图展现了相干类在零碎中的地位及其关系,便于后续追踪过程中查看。
OneFlow里定义了3个Stream类、2个Device类,后续剖析过程中留神防止混同。
让指令在虚拟机里执行
上一篇提到,在Interpret中,最终会结构一个lambda表达式让PhysicalRun
执行。
把传给PhysicalRun
的lambda表达式代入替换一下,理论执行如下代码:
vm::InstructionMsgList instruction_list; InstructionsBuilder instructions_builder(std::make_shared<vm::PhysicalIdGenerator>(), &instruction_list); // kernel等参数都由lambda绑定自Interpret instructions_builder.LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects, ctx, stream); JUST(vm::Run(instructions_builder.mut_instruction_list())); return Maybe<void>::Ok();
依据op_type获取kernel,结构虚拟机指令
LocalCallOpKernel这个函数很重要,这外面结构的变量在后续流程中都有重要作用。函数会结构一个InstructionMsg对象并放到列表中。
所谓指令,应该是OneFlow外部比拟细粒度的操作,而不是硬件指令。一个Op可能被转化为一个(或多个?)指令,交给调度引擎执行。
从类关系图也容易看出,指令是蕴含kernel和op conf信息的。
整个执行流程如下所示(始终到虚拟机接管指令):
LocalCallOpKernel函数中的instruction_name在CPU设施上就是"cpu.LocalCallOpKernel"。
LocalCallOpKernelPhyInstrOperand类型的对象phy_instr_operand,这个对象在Init时会调用ChooseOpKernel,从UserOpRegistryMgr获取OpKernelRegistryResult并调用它的create_fn成员函数,获取理论的kernel,对于relu来说就是ReluKernel。
而后创立InstructionMsg类型的变量instruction。这个变量是由intrusive::make_shared生成的(不是std::make_shared
)。这是OneFlow本人的援用计数实现。初始化调用的是InstructionMsg::\_\_Init\_\_办法。这个对象在初始化时,很重要的一个步骤是设置instr_type_id,其instruction_type的类型在CPU下就是CpuLocalCallOpKernelInstructionType。
同时会设置InstructionMsg
的stream。传入LocalCallOpKernel
的stream参数是oneflow::Stream
,而InstructionMsg
的stream是vm::Stream
。对于relu,指令stream最终来自GetDeviceStream,其中的Stream数组是在VM初始化时设置的,目前还不分明这个Stream数组的具体逻辑,不过能够确定其StreamType的类型是CpuStreamType。
从下面的类关系图能够看到,InstrTypeId类型涵盖了设施类型和[指令类型](https://github.com/Oneflow-In...
)。最终会在LookupInstrTypeId办法中设置InstrTypeId的值。上面须要找到InstrTypeId4InstructionName函数中的动态map是在哪里注册的。搜寻代码容易发现调用依赖关系如下:
- 在RegisterInstrTypeId函数中对map做了批改
- RegisterInstructionType调用RegisterInstrTypeId
- CpuLocalCallOpKernelInstructionType注册时调用RegisterInstructionType
注册宏开展后执行如下语句:
vm::RegisterInstructionType<CpuLocalCallOpKernelInstructionType>("cpu.LocalCallOpKernel");template<typename T>void RegisterInstructionType(const std::string& instr_type_name) { RegisterInstrTypeId<T>(instr_type_name, StaticGlobalStreamType<typename T::stream_type>());}
注册的key就是后面看到的instruction_name的值,value来自StaticGlobalStreamType
返回的动态变量。CpuLocalCallOpKernelInstructionType用于辨别StreamType,理论计算逻辑在LocalCallOpKernelInstructionType中。前面会看到,执行kernel计算时会调用这个类的办法。
虚拟机的初始化
在持续进入虚拟机之前,先看看虚拟机的初始化过程。虚拟机是OneFlow的执行引擎,VirtualMachine负责线程调度,具体任务交给VirtualMachineEngine执行。通过相似生产-生产的机制解决指令的执行。import oneflow
时在EnvGlobalObjectsScope::Init
中初始化虚拟机实例。具体过程如下:
在MakeVmDesc中,会把之前通过RegisterInstructionType
注册的InstrTypeId::stream_type_
都存到一个set中。再调用MakeStreamDesc结构StreamDesc对象,StreamDesc在结构时会设置stream_type(来自StaticGlobalStreamType保障指针惟一),对于relu来说就是CpuStreamType。最初将StreamDesc放到vm_desc.stream_type2desc中。这样,RegisterInstructionType
和VM中的StreamType指针是统一的。
在VirtualMachineEngine初始化时,依据StreamDesc顺次创立StreamRtDesc、ThreadCtx、vm::Stream,其中StreamType也是始终从StreamDesc传递下来。
虚拟机的调度机制
深度学习的Job能够视为一个有向无环图(DAG),算子/指令是图中的节点,节点是有依赖关系的。虚拟机负责保护若干个指令队列,以及指令在这些队列之间的状态转换。不同队列的指令有不同的依赖状态,比方刚收到期待调度、期待上游执行结束、能够被调度执行等。
指令结构结束后,调用Run交给虚拟机执行指令。在VirtualMachineEngine::Receive中,只是把指令列表放到pending_msg队列中。
指令的状态转换还没搞清楚,猜想大抵是这样的:
InstructionMsgList -> pending_msg
- Receive
pending_msg -> local_pending_msg
- Schedule
local_pending_msg -> ready_instruction
- HandleLocalPending
- GetRewritedPendingInstructionsByWindowSize
- MakeInstructions
ready_instruction -> Run
- Schedule
- DispatchAndPrescheduleInstructions
- DispatchInstruction
须要留神的是,Receive时收到的元素类型是InstructionMsg,ready_instruction的元素类型是Instruction,这个转换是在MakeInstructions内实现的。
指令调度与执行在逻辑上的调用程序如下:
追踪图中MakeInstructions的调用程序能够晓得,Instruction中的Stream和InstructionMsg中的指向同一个vm::Stream对象。
定位到具体的OpKernel
从上述状态转换来看,指令最终是在DispatchInstruction函数中执行的。这个函数执行指令的外围逻辑能够用如下伪码示意:
instruction->mut_stream()->stream->stream_type().Run(instruction);
依据下面InstructionMsg初始化的探讨,这里的StreamType就是CpuStreamType;instr_type_id.instruction_type的类型就是CpuLocalCallOpKernelInstructionType。这样就容易列出调用程序如下:
依据之前探讨的phy_instr_operand初始化的状况,OpKernelCompute中获取的user_opkernel就是ReluKernel,通过父类OpKernel的Compute办法进入ReluKernel::Compute。
在NewPrimitive中,须要搞清楚factory具体是什么类型。一路追踪到AutoRegistrationFactory,这里又是一个注册机制。然而用到REGISTER_CLASS的中央太多,一时仿佛没有脉络。
回头看NewReluPrimitive,这里指定的工厂类型是ElementwiseUnaryFactory。这是一个抽象类,搜寻一下容易发现它的CPU版本的子类ElementwiseUnaryImpl,其New办法定义了对各种数据类型的relu实现。这个源文件中还调用了宏REGISTER_PRIMITIVE_FACTORY。宏开展后的相干代码如下:
static AutoRegistrationFactory<DeviceType, ElementwiseUnaryFactory> ::RawRegisterType<ElementwiseUnaryFactoryImpl> g_registry_var0(DeviceType::kCPU);std::unique_ptr<ElementwiseUnary> New(UnaryOp unary_op, DataType src_type, DataType dst_dtype) override { static const std::map<std::tuple<UnaryOp, DataType, DataType>, std::function<std::unique_ptr<ElementwiseUnary>()>> new_elementwise_unary_handle { // ... { std::make_tuple((UnaryOp::kRelu), DataType::kFloat, DataType::kFloat), NewElementwiseUnary<(UnaryOp::kRelu), float, float> }, { std::make_tuple((UnaryOp::kRelu), DataType::kDouble, DataType::kDouble), NewElementwiseUnary<(UnaryOp::kRelu), double, double> }, // ... }; const auto it = new_elementwise_unary_handle.find( std::make_tuple(unary_op, src_type, dst_dtype)); if (it != new_elementwise_unary_handle.end()) { return it->second(); } else { return nullptr; }}
从以上代码容易看出,NewPrimitive返回的工厂类型是ElementwiseUnaryFactoryImpl。ReluKernel::Compute中的primitive类型是ElementwiseUnaryImpl。依据模版参数推断,其Launch办法中理论调用UnaryFunctor进行计算,在这里实现了relu的计算逻辑。
ReluKernel能够看作Kernel层对外的接口,由它依据context信息将工作转发给具体设施的计算函数。
小结
至此,Op执行相干的流程算是大体串了一遍。一句flow.relu()
前面会波及这么多内容。但这里其实也只关注了骨干逻辑,疏忽了两头大量的细节。
流程的梳理只是第一步,还须要从中演绎总结一些概念和概念之间的关系,再联合公开材料反推印证设计理念的落地实现。
不过目前对代码和设计的理解还很浮浅,上面的内容纯属大胆猜想。
Op执行的宏观脉络
从下面的类关系图登程,以外围类为节点,也能看出Op执行流程的宏观脉络。整个流程大体在上面这些角色之间流转:
- ReluFunctor
- UserOpExpr
- StatefulLocalOpKernel
- PhyInstrOperand
- InstructionMsg
- vm::Stream
用户结构的数据都会有设施属性,比方tensor是在CPU还是在GPU上计算。数据所在的设施信息封装在oneflow::Stream
类中。
UserOpExpr为每个oneflow::Stream
缓存一个StatefulLocalOpKernel。
StatefulLocalOpKernel向下能够依据UserOpRegistryMgr注册信息构建OpKernel,向上与Interpreter构建的PhyInstrOperand和指令关联。而指令也能够据此向下找到具体的Kernel执行计算。
Stream
OneFlow中,硬件资源,包含CPU、GPU和网络等都被形象成工作队列,对立把这样的队列称为stream。
OneFlow中有3个Stream类,别离是:
- oneflow::Stream: tensor数据的设施信息用这个类示意。
- vm::Stream: 更像是负责虚拟机的计算资源管理和调度。
- ep::Stream: 示意具体的计算设施,有CPU和GPU等不同类型的子类实现。比方可能会提供OneDnn等反对。
那么,oneflow::Stream和ep::Stream为什么要分2个类呢?猜想一下,比方跨设施的运算、数据搬运等,数据输出与理论计算的设施可能会不一样。从下面剖析的执行流程看,将两个Stream串起来的应该是来自inputs的device_id。不过具体细节设计vm初始化时的设施解决,目前还没搞清楚。
UserOpExpr
UserOpExpr示意一个具体的算子。其实UserOp只是Op中的一种。下图展现了不同Op的继承关系。能够看到tensor转换等也都视为Op。
参考资料
- 从OpExprInterpreter到OpKernel
- 动静调度的“咒骂”| 原有深度学习框架的缺点③
- OneFlow