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

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

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