关于深度学习:OneFlow源码阅读3Op指令在虚拟机中的执行

24次阅读

共计 6299 个字符,预计需要花费 16 分钟才能阅读完成。

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

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

正文完
 0