OneFlow 提供简洁高效易扩展的硬件抽象层 EP（Execution Provider），以应对适配不同硬件的复杂性。 引入硬件抽象层之后，用户无需关注底层硬件和框架的具体实现细节，框架的各个模块无需改动便可以适配新的硬件设备，同时，用户只需按照硬件抽象接口的约定和硬件设备的实际情况，实现一系列接口，便可以完成硬件的适配工作。
EP 还定义了一组基础计算接口 Primitive，基于 Primitive 接口重新实现了 Kernel。 相比 EP 提供的运行时接口，Primitive 提供的接口更加灵活，不同接口之间相互独立，每一个接口表示了某种硬件设备可以提供的特定的计算能力。
ep 模块主要包括两部分。一部分是之前讨论的设备管理，根据用户提供的信息能获取设备实例，将设备抽象出 Stream、Event、内存管理等接口。
另一部分就是基础计算接口 Primitive。这里只简要介绍一下 Primitive 的概念，包含哪些内容。不会涉及具体计算的设计和实现。
1
Primitive 是什么？
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Backward |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
||
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2
部分计算接口的说明
2.1 ElementwiseUnary
2.1.1 relu kernel 的执行过程
auto primitive_factory_func_ = [](user_op::KernelComputeContext* ctx) {const user_op::TensorDesc* src = ctx->TensorDesc4ArgNameAndIndex("x", 0);const user_op::TensorDesc* dst = ctx->TensorDesc4ArgNameAndIndex("y", 0);return ep::primitive::NewPrimitive<ep::primitive::ElementwiseUnaryFactory>(ctx->device_type(), ep::primitive::UnaryOp::kRelu, src->data_type(),dst->data_type());};OpKernel* ptr = new UnaryPrimitiveKernel("y", "x", primitive_factory_func_);
-
调用 primitive_factory_func_ 获取一个 primitive 实例。 -
NewPrimitive -
调用 NewObjUniquePtr 获取 ElementwiseUnaryFactoryImpl 实例（ CPU ， CUDA ）。 -
调用 ElementwiseUnaryFactoryImpl::New 返回 ElementwiseUnaryImpl 实例（ CPU , CUDA ）。 -
调用 primitive->Launch 执行计算。
2.1.2 ElementwiseUnary 支持哪些操作？
static const std::map<std::tuple<UnaryOp, DataType, DataType>,std::function<std::unique_ptr<ElementwiseUnary>(Scalar, Scalar)>>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>},{std::make_tuple((UnaryOp::kElu), DataType::kFloat, DataType::kFloat), NewElementwiseUnary<(UnaryOp::kElu), float, float>},{std::make_tuple((UnaryOp::kLogicalNot), DataType::kDouble, DataType::kBool), NewElementwiseUnary<(UnaryOp::kLogicalNot), double, bool>},// ......};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(attr0, attr1);} else {return nullptr;}
2.1.3 ElementwiseUnaryImpl::Launch 的实现
-
各设备通用的 UnaryFunctor 实现 。其中包括 relu 的实现 。 -
CPU 的 UnaryFunctor 实现 。通过 cpu_stream->ParallelFor 并行加速。 -
CUDA 的 UnaryFunctor 实现 。后续通过 cuda::elementwise::Unary 调用设备计算。
2.2 BroadcastElementwiseBinary
-
broadcast_elementwise_binary_activation_grad.cu -
broadcast_elementwise_binary_comparision.cu -
broadcast_elementwise_binary_logical.cu -
broadcast_elementwise_binary_math.cu
nvcc -DWITH_CUDA \-E -std=c++17 \-I. -Ibuild \-Ibuild/oneflow/ir/llvm_monorepo-src/llvm/include \-Ibuild/oneflow/ir/llvm_monorepo-build/include \-Ibuild/half/src/half/include \-Ibuild/_deps/glog-src/src -Ibuild/_deps/glog-build \-Ibuild/protobuf/src/protobuf/src \oneflow/core/ep/cuda/primitive/broadcast_elementwise_binary_math.cu > math.cpp
3
UserOp、Kernel 与 Primitive 的关系
3.1 UserOp 与 Kernel 是一对多的关系
REGISTER_USER_KERNEL("max_pool_2d").SetCreateFn<MaxPool2dKernel<device, dtype().SetIsMatchedHob((user_op::HobDeviceType() == device)&& (user_op::HobDataType("x", 0) == GetDataType<dtype>::value));
-
kernel_reg_val = UserOpRegistryMgr::Get().GetOpKernelRegistryResult(...) -
通过 reg_val.is_matched_hob->get(ctx) 判断 Kernel 是否匹配 -
如果 没有匹配 会报错。如果 多于一个匹配会报警 -
kernel = kernel_reg_val->create_fn()
3.2 IsMatchedHob 到底是啥？
using IsMatchedHob = std::shared_ptr<hob::BaseExpr<user_op::KernelRegContext, bool;
3.2.1 布尔表达式的析构函数
template<typename T>OpKernelRegistry& SetIsMatchedHob(const T& hob) {result_.is_matched_hob = std::make_shared<T>(hob);return *this;}
参考资料
试用OneFlow: github.com/Oneflow-Inc/oneflow/
本文分享自微信公众号 - OneFlow（OneFlowTechnology）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。
