共计 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