下图展现了相干类在零碎中的地位及其关系,便于后续追踪过程中查看。

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