td目标源码(Td目标禁止)

发布时间： 2023-03-27 12:37:52

　　

撰文｜姚迟、郑泽康

　　

本文将以开发一个 leaky_relu（准确说是 leaky_relu_yzh op，因为 master 分支的 leaky_relu 组合了其它知识点）为例介绍如何在 OneFlow 中新增算子（https://github.com/Oneflow-Inc/oneflow/pull/8350）。

　　


　　

1

　　

背景

op 与 kernel

　　

op 与 kernel 是两个有关联的概念。op 是逻辑上的算子，包含 OneFlow Compiler 在构建计算图时所需要的必要信息，如输入、输出形状，哪些张量需要自动求导等信息。有了 op 中的信息，OneFlow Compiler 就可以构建计算图并依据计算图做资源申请、构建等操作（如根据张量的输入输出大小申请内存）， 但是 op 中不包含具体的处理数据的逻辑。

　　

在真正需要处理数据时，OneFlow Runtime 会启动 kernel 完成计算，所以 kernel 中包含了具体处理数据的逻辑。对于一个逻辑上的 op，OneFlow Runtime 会根据数据类型、硬件设备（比如是 CPU 还是 CUDA）的具体情况，选择启动不同的 kernel。

　　

OneFlow 中的系统 op 与 user op

　　

在 OneFlow 系统中存在两类算子（op）：系统 op 和 user op。

　　

系统 op 定义在：oneflow/core/operator/ 目录， 对应的 kernel 实现在：oneflow/core/kernel 目录。系统 op 是对构图、流水等系统性能较为关键的一些 op。

　　

除极少数 op 属于系统 op 外，大多数 op 都是 user op，这些 user op 和用户模型业务逻辑相关。OneFlow user op 的定义及 kernel 实现分别在 oneflow/user/ops 和 oneflow/user/kernels 目录下。

　　

目前 OneFlow 已实现了丰富的算子库，但是当已有的算子库无法满足搭建模型的需求时，就需要新增算子。本文介绍的新增算子指的是新增 user op。

　　

ODS 与 TableGen

　　

TableGen（https://llvm.org/docs/TableGen/index.html） 是一个代码生成工具，简单而言，它读取并解析一个 .td 格式（语法接近 C++ 模板）的文件，然后交给 TableGen 后端

　　

（https://llvm.org/docs/TableGen/BackEnds.html）生成另外格式的语言。

　　

MLIR 基于 TableGen 制定了一套算子定义规范ODS（https://mlir.llvm.org/docs/OpDefinitions/）以及对应的后端 OpDefinitionsGen（https://github.com/llvm/llvm-project/blob/main/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp。）

　　

OneFlow 在 ODS 的基础上，实现了 TableGen OneFlow 后端（https://github.com/Oneflow-Inc/oneflow/tree/master/tools/oneflow-tblgen），并使用它来定义 OneFlow user op。

　　

因此，OneFlow 的 user op 定义写在 OneFlowUserOps.td 文件中。

　　


　　

2

　　

开发 op

在 OneFlow 中开发一个新的 user op，主要分为以下4步：

　　

  1. 定义 op
  2. 实现 kernel 计算逻辑
  3. 导出 functional 接口
  4. 实现用于求导的反向逻辑

定义 op

　　

定义 op 指的是，对 op 的名称，op 的输入、输出数据类型和 op 的属性进行声明。OneFlow 遵循 MLIR 的 ODS（Operation Definition Specification）（
https://mlir.llvm.org/docs/OpDefinitions/） 实现了自己的 MLIR OneFlow Dialect。在算子定义方面，这样做的好处是，各种推导函数和序列化/反序列化的接口都可以委托给 ODS，降低了人工手写出错的概率，后续优化、格式转化等流程可以更灵活。

　　

定义一个 OneFlow user op，主要包括 5 个部分，分别是：

　　

  • op class
  • 输入 input
  • 输出 output
  • 属性 attrs
  • 导出并实现推导接口

op class

　　

可以在
oneflow/ir/include/OneFlow/OneFlowUserOps.td 查看 op 定义的源码。

　　

以 def 关键字开头定义一个 op，该 op 继承 OneFlow_BaseOp，同时指定 OneFlow_BaseOp 的模版参数。模版参数依次为 op type name、Trait （
https://mlir.llvm.org/docs/Traits/）列表。

　　

defOneFlow_LeakyReluYZHOp : OneFlow_BaseOp<"leaky_relu_yzh", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
//...
}

其中 "leaky_relu_yzh" 是指定的 op type name。每个 op 都需要指定一个全局唯一的 op type name 作为全局标识符。

　　

第二个模板参数是一个 list（[...]），其中的每一项都是一个 Trait，OneFlow 中常用的有：

　　

  • NoSideEffect 表示该算子无副作用（即不会改变内存、网络、管道、磁盘等的系统状态），这个特性可以指导某些优化操作
  • NoGrad 表示该算子在数学上没有梯度（不可导）
  • CpuOnly 表示该算子只支持在 CPU 设备上执行
  • SupportNonContiguous 表示该算子是否支持 NonContiguous 张量（关于 Contiguous Tensor 的概念，可以参考 PyTorch Internals 中的相关内容 ）

输入 input 与输出 output

　　

通过重写 input 域来定义 op 的输入，比如

　　

// 一个输入 x
let input = (ins
 OneFlow_Tensor:$x
);

定义了一个输入张量 x。输入的格式为 输入类型:$name。

　　

输入类型目前包括：

　　

  • OneFlow_Tensor
  • Variadic<OneFlow_Tensor>：指可变 tensor，比如 concat op，支持 concat 可变个数的 tensor。
  • Optional<OneFlow_Tensor>：表示这个 tensor 是可选的，既可以有也可以没有，比如 conv op 中的 add_output。

一个 op 也可以定义多个输入，比如：

　　

// 两个输入：a, b
 let input = (ins
 OneFlow_Tensor:$a,
 OneFlow_Tensor:$b
 );

通过重写 output 域来定义 op 的输出，比如下面定义了 2 个输出张量：

　　

let output = (outs
 OneFlow_Tensor:$out0,
 OneFlow_Tensor:$out1
);

属性 attrs

　　

通过重写 attrs 域定义 op 的属性，比如定义 dropout （
https://oneflow.readthedocs.io/en/master/functional.html#
oneflow.nn.functional.dropout）中的 rate 属性：

　　

let attrs = (ins
 DefaultValuedAttr<F32Attr, "0.">:$rate
 );

它表示名为 $rate 的类型是 F32Attr，默认值是 0.。这里也可以不指定默认值：

　　

let attrs = (ins
 F32Attr:$rate
 );

I32Attr、F32Attr、BoolAttr、StrAttr、I32ArrayAttr 等常见基础数据类型定义在 OpBase.td

　　

（https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/OpBase.td#L1077-L1086）中。

　　

OneFlow 自定义数据类型，如 ShapeAttr、DTArrayAttr 等定义在 OneFlowBase.td

　　

（https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/ir/include/OneFlow/OneFlowBase.td#L27-L35）中。

　　

导出并实现推导接口

　　

还有一些其它域，用于指定是否生成对应的接口。这些接口往往是构建计算图过程中的推导接口。

　　

比如 shape 推导（根据输入的 shape 推导输出的推导）、data type 推导、SBP 推导等。

　　

OneFlow-TableGen 仅负责生成这些函数的接口，开发者需要在其自动生成的 cpp 文件中实现这些接口。默认情况不会生成下列任何接口，开发者需要显式指定需要生成哪些接口。

　　

let has_check_fn = 1; // 生成属性检查接口
 let has_logical_tensor_desc_infer_fn = 1; // 生成 logical shape 推导接口
 let has_physical_tensor_desc_infer_fn = 1; // 生成 physical shape 推导接口
 let has_get_sbp_fn = 1; // 生成 get sbp 接口
 let has_sbp_signature_infer_fn = 1; // 生成 sbp signature 推导接口，未来会移除，推荐使用 has_nd_sbp_infer_fn
 let has_data_type_infer_fn = 1; // 生成 data type 推导接口
 let has_device_and_stream_infer_fn = 1; // 生成 device 推导接口
 let has_input_arg_modify_fn = 1; // 生成输入 modify 接口，比如设置 is_mutable、requires_grad（用于Lazy）等
 let has_output_arg_modify_fn = 1; // 生成输出 modify 接口，比如设置 is_mutable、requires_grad（用于Lazy）等
 let has_output_blob_time_shape_infer_fn = 1; // 生成输出 time shape 推导接口
 let has_nd_sbp_infer_fn = 1; // 生成 nd sbp 推导接口

一般常用的是下面几个：

　　

let has_logical_tensor_desc_infer_fn = 1;
 let has_physical_tensor_desc_infer_fn = 1;
 let has_data_type_infer_fn = 1;
 let has_get_sbp_fn = 1;

了解完上面这些概念和用法后，可以开始修改
oneflow/ir/include/OneFlow/OneFlowUserOps.td文件。

　　

leaky_relu_yzh op 完整的定义见 这里（
https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/ir/include/OneFlow/OneFlowUserOps.td#L8418-L8433）。

　　

在 OneFlowUserOps.td 中新增Op定义之后，重新 make 后会自动在 build 目录下的 oneflow/core/framework/ 目录下生成文件以下几个文件：

　　

  • op_generated.h：由解析 .td 文件生成的 op C++ 类
  • op_generated.cpp：由解析 .td 文件生成的 op 注册代码（包含调用 REGISTER_USER_OP 宏的代码）

之后需要做的就是在 oneflow/user/ops （https://github.com/Oneflow-Inc/oneflow/tree/master/oneflow/user/ops）目录下新加一个 cpp 文件，用于实现 op 的接口。

　　

比如 leaky_relu_yzh 对应的文在 oneflow/user/ops/leaky_relu_yzh_op.cpp（https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/user/ops/leaky_relu_yzh_op.cpp#L21-L79），实现了推导逻辑张量、推导物理张量、推导 SBP 信息以及推导输出数据类型各接口。

　　

实现 Kernel 逻辑

　　

op 的计算支持多种设备（如 CPU、GPU、DCU 等），所以要分别实现计算逻辑。

　　

相关代码：

　　

  • Leaky ReLU CPU Kernel
  • （https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/user/kernels/leaky_relu_yzh_kernel.cpp）
  • Leaky ReLU GPU KernelCPU
  • （https://github.com/Oneflow-Inc/oneflow/blob/7ab4b0f08c86a6f8af08b44daa510725942288fb/oneflow/user/kernels/leaky_relu_yzh_kernel.cu）

计算逻辑

　　

template<typename T>
class CpuLeakyReluYZHKernel final : public user_op::OpKernel {
 public:
 CpuLeakyReluYZHKernel() = default;
 ~CpuLeakyReluYZHKernel() = default;
 private:
 void Compute(user_op::KernelComputeContext* ctx) constoverride 
 const user_op::Tensor* x = ctx->Tensor4ArgNameAndIndex("x", 0);
 user_op::Tensor* y = ctx->Tensor4ArgNameAndIndex("y", 0);
 const int32_t elem_cnt = x->shape().elem_cnt();
 const T* x_ptr = x->dptr<T>();
 T* y_ptr = y->mut_dptr<T>();
 const auto alpha = ctx->Attr<float>("alpha");
 FOR_RANGE(int32_t, i, 0, elem_cnt) { y_ptr[i] = x_ptr[i] > 0 ? x_ptr[i] : alpha * x_ptr[i]; }
 }
 bool AlwaysComputeWhenAllOutputsEmpty() constoverride { returnfalse; }
};

在 OneFlow 中实现 kernel， 必须定义一个继承自
oneflow::user_op::OpKernel 的类，并重写其中的虚函数。在以上代码。