OneFlow源码阅读3:Op指令在虚拟机中的执行

2022年05月15日 阅读数:6
这篇文章主要向大家介绍OneFlow源码阅读3:Op指令在虚拟机中的执行,主要内容包括基础应用、实用技巧、原理机制等方面,希望对大家有所帮助。

下图展现了相关类在系统中的位置及其关系,便于后续追踪过程当中查看。git

OneFlow里定义了3个Stream类、2个Device类,后续分析过程当中注意避免混淆。
github

让指令在虚拟机里执行

上一篇提到,在Interpret中,最终会构造一个lambda表达式让PhysicalRun执行。segmentfault

把传给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信息的。框架

整个执行流程以下所示(一直到虚拟机接收指令):
Interpretide

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是在哪里注册的。搜索代码容易发现调用依赖关系以下:

注册宏展开后执行以下语句:

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。这样就容易列出调用顺序以下:
op到kernel

根据以前讨论的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。
OpExpr类关系

参考资料