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

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

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

1 指令在虚拟机里执行

上一篇提到，在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();

1.1 根据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计算时会调用这个类的方法。

2 虚拟机的初始化

在继续进入虚拟机之前，先看看虚拟机的初始化过程。虚拟机是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传递下来。

2.1 虚拟机的调度机制

深度学习的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对象。

3 定位到具体的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信息将任务转发给具体设备的计算函数。

4 小结

至此，Op执行相关的流程算是大体串了一遍。一句flow.relu()后面会涉及这么多内容。但这里其实也只关注了主干逻辑，忽略了中间大量的细节。

流程的梳理只是第一步，还需要从中归纳总结一些概念和概念之间的关系，再结合公开资料反推印证设计理念的落地实现。

不过目前对代码和设计的了解还很肤浅，下面的内容纯属大胆猜测。

4.1 Op执行的宏观脉络

从上面的类关系图出发，以核心类为节点，也能看出Op执行流程的宏观脉络。整个流程大体在下面这些角色之间流转：

• ReluFunctor
• UserOpExpr
• StatefulLocalOpKernel
• PhyInstrOperand
• InstructionMsg
• vm::Stream

用户构造的数据都会有设备属性，比如tensor是在CPU还是在GPU上计算。数据所在的设备信息封装在oneflow::Stream类中。

UserOpExpr为每个oneflow::Stream缓存一个StatefulLocalOpKernel。

StatefulLocalOpKernel向下可以根据UserOpRegistryMgr注册信息构建OpKernel，向上与Interpreter构建的PhyInstrOperand和指令关联。而指令也可以据此向下找到具体的Kernel执行计算。

4.2 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初始化时的设备处理，目前还没搞清楚。

4.3 UserOpExpr

UserOpExpr表示一个具体的算子。其实UserOp只是Op中的一种。下图展示了不同Op的继承关系。可以看到tensor转换等也都视为Op。

参考资料

• 从OpExprInterpreter到OpKernel
• 动态调度的“诅咒”| 原有深度学习框架的缺陷③
• OneFlow