关于深度学习:OneFlow-源码阅读-11基础计算接口-Primitive

OneFlow v0.8.0 公布文档中的第 5 节对框架的“多设施适配”作了阐明，原文摘录如下：

OneFlow 提供简洁高效易扩大的硬件形象层 EP（Execution Provider），以应答适配不同硬件的复杂性。引入硬件形象层之后，用户无需关注底层硬件和框架的具体实现细节，框架的各个模块无需改变便能够适配新的硬件设施，同时，用户只需依照硬件形象接口的约定和硬件设施的理论状况，实现一系列接口，便能够实现硬件的适配工作。
EP 还定义了一组根底计算接口 Primitive，基于 Primitive 接口从新实现了 Kernel。相比 EP 提供的运行时接口，Primitive 提供的接口更加灵便，不同接口之间互相独立，每一个接口示意了某种硬件设施能够提供的特定的计算能力。

ep 模块次要包含两局部。一部分是之前探讨的设施治理，依据用户提供的信息能获取设施实例，将设施形象出 Stream、Event、内存治理等接口。

另一部分就是根底计算接口 Primitive。这个笔记只粗略讨论一下 Primitive 的概念，大抵是什么样子的，蕴含哪些内容；但不会波及具体计算的设计和实现。

1 primitive 是什么？

粗略地说，根底计算接口是指 primitive 目录下定义的二十来个根底计算接口类。它们都是 Primitive 的子类。这些接口类型通常只申明一个 Launch 办法，理论反对哪些计算是由针对具体设施的实现决定的。

各根底计算接口如下表所示：

2 局部计算接口的阐明

2.1 ElementwiseUnary

2.1.1 relu kernel 的执行过程

relu kernel 就是通过 ElementwiseUnary 执行计算的。注册 relu kernel 的 SetCreateFn 函数执行相似如下代码的操作。UnaryPrimitiveKernel 结构时会保留 primitive_factory_func_。

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_);

在调用 UnaryPrimitiveKernel::Compute 执行 kernel 计算时，执行如下操作：

• 调用 primitive_factory_func_ 获取一个 primitive 实例。

  • NewPrimitive

    • 调用 NewObjUniquePtr 获取 ElementwiseUnaryFactoryImpl 实例（CPU，CUDA）。
    • 调用 ElementwiseUnaryFactoryImpl::New 返回 ElementwiseUnaryImpl 实例（CPU, CUDA）。
• 调用 primitive->Launch 执行计算。

上述类之间的关系如下：

2.1.2 ElementwiseUnary 反对哪些操作？

ElementwiseUnaryFactoryImpl::New 中的宏开展后，代码如下。依据 UnaryOp 的操作类别、数据类型查到 New 函数，传递对应的模版参数给 New 函数并创立 ElementwiseUnaryImpl 实例。
ElementwiseUnary 在 CPU 环境反对的 < 操作, 数据类型 > 组合都在这个 map 中注册。这个应该就是 “惯例”意义上的“Primitive 接口” 的一部分（反对哪些操作、数据类型等），操作的输出参数由 Launch 函数的接口决定。

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 的实现

Primitive 不同子类的 Launch 办法，其实现形式和输出参数各不一样。ElementwiseUnaryImpl::Launch 通过 primitive::UnaryFunctor 实现计算逻辑（CPU，CUDA）。

primitive::UnaryFunctor 是一个模版类，其特化版本散布在如下文件：

• 各设施通用的 UnaryFunctor 实现。其中包含 relu 的实现。
• CPU 的 UnaryFunctor 实现。通过 cpu_stream->ParallelFor 并行减速。
• CUDA 的 UnaryFunctor 实现。后续通过 cuda::elementwise::Unary 调用设施计算。

2.2 BroadcastElementwiseBinary

BroadcastElementwiseBinary 也定义了 CUDA 的工厂实现。New 函数的 map 中定义了 CUDA 下反对的所有操作组合，每个都是一个 NewBroadcastElementwiseBinary 模版函数的特化实例的援用。这些模版函数的特化定义在上面几个文件中：

• broadcast_elementwise_binary_activation_grad.cu
• broadcast_elementwise_binary_comparision.cu
• broadcast_elementwise_binary_logical.cu
• broadcast_elementwise_binary_math.cu

这些文件中的宏能够用如下命令开展，必须指定 WITH_CUDA 能力失常开展宏。

nvcc -DWITH_CUDA \
  -E -std=c++14 \
  -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 Primitive 仿佛并未笼罩全副的 Kernel

绝大部份 Kernel 都用 Primitive 实现计算逻辑。然而也有局部 Kernel 没用 Primitive，而是间接调用设施办法，比方 conv。

3.2 UserOp 与 Kernel 是一对多的关系

之前看过的代码，UserOp 通常只有一个 Kernel，Kernel 不辨别设施、通过 Primitive 适配不同的设施计算。但也有例外。
通过 conv kernel 能够看到，CPU 和 CUDA 注册了同名的 kernel。认真看 UserOpRegistryMgr::op_kernel_reg_result_ 的 value 类型是 vector。所以 UserOp 与 Kernel 是一对多的关系。通过 OpKernelRegistryResult::is_matched_hob 筛选出匹配的 kernel。
以 max_pool_2d 为例，其 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 计算的筹备阶段，在 StatefulOpKernel::ChooseOpKernel 中相干调用如下：

• kernel_reg_val = UserOpRegistryMgr::Get().GetOpKernelRegistryResult(…)

  • 通过 reg_val.is_matched_hob->get(ctx) 判断 Kernel 是否匹配
  • 如果没有匹配会报错。如果多于一个匹配会报警（之前是报错。v0.9.0 仿佛引入了 kernel 优先级的概念？）
• kernel = kernel_reg_val->create_fn()

3.3 IsMatchedHob 到底是啥？

is_matched_hob 的类型是 IsMatchedHob：

using IsMatchedHob = std::shared_ptr<hob::BaseExpr<user_op::KernelRegContext, bool>>;

(user_op::HobDeviceType() == device) && (user_op::HobDataType("x", 0) == GetDataType<dtype>::value) 并不是一个一般的 bool 表达式，而是一个相似下图的高阶表达式：

HobDeviceType() 返回的类型是 Custom，它是 Expr 的子类，其 ValueT 是 DeviceType。DEFINE_BINARY_FUNCTOR 宏定义了一个重载 Expr 的 == 运算符的函数，第一个参数类型是 Expr（也就是 Custom），第二个参数类型是 Custom::ValueT，也就是 DeviceType，返回的 BoolFunctor 继承自 BoolExpr，也是 Expr 的子类。相似的，也通过宏定义了 And 运算符的重载。这样就形成了如上图所示的高阶 bool 表达式。BoolFunctor::get 函数在运行时依据 context 动静计算表达式的值。比方 normalization 用来辨别是训练还是推理。

各类型关系如下：

3.3.1 布尔表达式的析构函数

BaseExpr 是上述这些 bool 表达式对象的基类。其析构函数不是 virtual 的。SetIsMatchedHob 的代码如下。调用时 T 的具体类型是确定的，make_shared 晓得如何正当开释，所以这个场景不会造成内存透露。框架应该只是在 kernel 注册相干环节应用了这些类型。

  template<typename T>
  OpKernelRegistry& SetIsMatchedHob(const T& hob) {result_.is_matched_hob = std::make_shared<T>(hob);
    return *this;
  }

4 参考资料

• oneflow v0.9.0