在线构建自定义算子
Windows Linux Android C++ 推理应用 高级
如何实现自定义算子
MindSpore Lite当前提供了一套南向的算子注册机制，如果用户想通过MindSpore Lite框架调度到自己的算子实现上，可参考本文。
实现自定义算子大概有以下几个步骤：
确定算子类型 ：分为通用算子与Custom算子。
算子实现：继承Kernel类实现自定义算子，并注册进MindSpore Lite。
算子InferShape：继承mindspore::kernel::KernelInteface实现自定义算子的InferShape能力，并注册进MindSpore Lite。
确定算子类型
查看mindspore/lite/schema/ops.fbs中的算子原型定义，确认要注册实现的算子原型是否在PrimitiveType中有定义，有定义的话则要注册的算子为通用算子，可以按照已有的IR直接实现算子与注册，否则即为Custom算子。
通用算子
整个算子的实现、注册、infershape等相关的代码可以参看代码仓里的样例。
通用算子实现
继承mindspore::kernel::Kernel，重载实现必要的接口。以自定义一个Add算子为例：
算子继承Kernel。
PreProcess()对内存进行了预分配。
Execute()对input进行了相加。
using mindspore::kernel::Kernel;
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_OK;
using mindspore::lite::RET_PARAM_INVALID;
class TestCustomAdd : public Kernel {
public:
TestCustomAdd(const std::vector<tensor::MSTensor *> &inputs, const std::vector<tensor::MSTensor *> &outputs,
const schema::Primitive *primitive, const lite::Context *ctx)
: Kernel(inputs, outputs, primitive, ctx) {}
int Prepare() override { return 0; }
int Execute() override;
int ReSize() { return 0; }
private:
int PreProcess() {
for (auto *output : outputs_) {
// malloc data for output tensor
auto data = output->MutableData();
if (data == nullptr) {
MS_LOG(ERROR) << "Get data failed";
return RET_ERROR;
}
}
return RET_OK;
}
};
int TestCustomAdd::Execute() {
if (inputs_.size() != 2) {
return RET_PARAM_INVALID;
}
PreProcess();
auto *in0 = static_cast<const float *>(inputs_[0].Data().get());
auto *in1 = static_cast<const float *>(inputs_[1].Data().get());
float *out = static_cast<float *>(outputs_[0].MutableData());
auto num = outputs_[0].ElementNum();
for (int i = 0; i < num; ++i) {
out[i] = in0[i] + in1[i];
}
return RET_OK;
}
通用算子注册
当前有提供现成的宏REGISTER_KERNEL可以进行算子注册，实现步骤如下：
函数TestCustomAddCreator用来创建Kernel。
通过宏REGISTER_KERNEL进行Kernel注册，这里产商假定为BuiltInTest。
using mindspore::schema::PrimitiveType_AddFusion;
std::shared_ptr<Kernel> TestCustomAddCreator(const std::vector<tensor::MSTensor *> &inputs,
const std::vector<tensor::MSTensor *> &outputs,
const schema::Primitive *primitive, const lite::Context *ctx) {
return std::make_shared<TestCustomAdd>(inputs, outputs, primitive, ctx);
}
const auto kFloat32 = DataType::kNumberTypeFloat32;
REGISTER_KERNEL(CPU, BuiltInTest, kFloat32, PrimitiveType_AddFusion, TestCustomAddCreator)
通用算子InferShape
继承KernelInterface后重载Infer函数，实现InferShape能力。实现步骤如下：
重载实现Infer函数，推导出output tensor的shape，format，data_type。
这里以自定义Add算子为例：
using mindspore::kernel::KernelInterface;
class TestCustomAddInfer : public KernelInterface {
public:
TestCustomAddInfer() = default;
~TestCustomAddInfer() = default;
Status Infer(std::vector<mindspore::MSTensor *> *inputs, std::vector<mindspore::MSTensor *> *outputs,
const schema::Primitive *primitive) override {
(*outputs)[0].SetFormat((*inputs)[0].format());
(*outputs)[0].SetDataType((*inputs)[0].DataType());
(*outputs)[0].SetShape((*inputs)[0].Shape());
return kSuccess;
}
};
通用算子InferShape注册
当前有提供现成的宏REGISTER_KERNEL_INTERFACE可以进行算子InferShape注册，步骤如下：
函数CustomAddInferCreator用来创建KernelInterface实例。
调用REGISTER_KERNEL_INTERFACE宏对通用算子InferShape进行注册，这里产商假定为BuiltInTest。
std::shared_ptr<KernelInterface> CustomAddInferCreator() { return std::make_shared<TestCustomAddInfer>(); }
REGISTER_KERNEL_INTERFACE(BuiltInTest, PrimitiveType_AddFusion, CustomAddInferCreator)
Custom算子
Custom算子的解析、创建、操作等相关的代码可以参看代码仓里的样例。
Custom算子定义
table Attribute {
name: string;
data: [ubyte];
}
table Custom {
type: string;
attr: [Attribute];
}
属性是以字典的形式进行存储：name解释了属性名，data里存储了属性内容的字节流。 type：Custom算子的类型。
Custom算子创建
通过转换工具Converter的Pass注册接口，可以注册用户自己的Pass，用以导出想要的算子结构。这里以AddN算子转为一个Custom算子为例：
设Custom算子存在”input_num”、”op_kind”属性。
通过自定义Pass子类，实现Custom算子的转换与创建。
注册自定义Pass类。
namespace mindspore::opt {
class Test2Fusion : public Pass {
public:
AnfNodePtr CreateCustomOp(const FuncGraphPtr func_graph, const CNodePtr cnode) {
if (func_graph == nullptr || cnode == nullptr) {
return nullptr;
}
auto primc = std::make_shared<ops::Custom>(); // 创建Primitive，存储算子属性
if (primc == nullptr) {
return nullptr;
}
primc->set_type("Custom_AddN"); // 设置Custom算子类型
std::map<std::string, std::vector<uint8_t>> custom_attrs;
std::string input_num = std::to_string(cnode->size() - 1);
std::vector<uint8_t> input_num_attr(input_num.begin(), input_num.end());
custom_attrs["input_num"] = input_num_attr;
std::string op_kind = "custom op";
std::vector<uint8_t> op_kind_attr(op_kind.begin(), op_kind.end());
custom_attrs["op_kind"] = op_kind_attr;
primc->set_attr(custom_attrs); // 设置Custom算子属性
auto inputs = cnode->inputs();
inputs.erase(inputs.begin());
auto custom_cnode = func_graph->NewCNode(primc, inputs); // 创建CNode节点
custom_cnode->set_fullname_with_scope(cnode->fullname_with_scope()); // 设置节点名
custom_cnode->set_abstract(cnode->abstract()->Clone()); // 设置算子输出的基本属性，存储于abstract中
return custom_cnode;
}
bool Run(const FuncGraphPtr &func_graph) override {
auto manager = Manage(func_graph, true); // 创建FuncGrap管理器
if (manager == nullptr) {
return false;
}
auto node_list = TopoSort(func_graph->get_return()); // 获取所有节点
for (auto &node : node_list) {
if (!utils::isa<CNode>(node)) {
continue;
}
if (!opt::CheckPrimitiveType(node, prim::kPrimAddN)) { // 判断当前节点是否为AddN算子
continue;
}
auto cnode = node->cast<CNodePtr>();
auto custom_cnode = CreateCustomOp(func_graph, cnode); // 创建Custom算子
if (custom_cnode == nullptr) {
return false;
}
manager->Replace(node, custom_cnode) // 通过管理器用新节点替换旧节点
}
return true;
}
};
REG_PASS(Test1Fusion, Test1Fusion) // 注册Test1Fusion
REG_PASS(Test2Fusion, Test2Fusion) // 注册Test2Fusion
std::vector<std::string> schedule = {"Test1Fusion", "Test2Fusion"};
REG_SCHEDULED_PASS(POSITION_BEGIN, schedule) // 设置外部Pass调度逻辑，在内置融合前运行外部Pass
} // namespace mindspore::opt
整个Custom算子的实现、注册、infershape等相关的代码可以参看代码仓里的样例。
Custom算子实现
Custom算子的实现整体流程与通用算子的实现是一致的，因为都是Kernel的具体子类。 如果自定义算子不是运行在CPU平台上，需要在运行结束时把结果重新拷回output tensor。这里以创建一个Add能力的Custom算子为例：
算子继承Kernel。
PreProcess()对内存进行了预分配。
Execute()对input进行了相加。
using mindspore::kernel::Kernel;
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_OK;
class TestCustomOp : public Kernel {
public:
TestCustomOp(const std::vector<tensor::MSTensor *> &inputs, const std::vector<tensor::MSTensor *> &outputs,
const schema::Primitive *primitive, const lite::Context *ctx)
: Kernel(inputs, outputs, primitive, ctx) {}
int Prepare() override { return 0; }
int Execute() override;
int ReSize() override { return 0; }
private:
int PreProcess() {
for (auto *output : outputs_) {
// malloc data for output tensor
auto data = output->MutableData();
if (data == nullptr) {
MS_LOG(ERROR) << "Get data failed";
return RET_ERROR;
}
}
return RET_OK;
}
int TestCustomOp::Execute() {
if (inputs_.size() != 2) {
return RET_PARAM_INVALID;
}
PreProcess();
GetAttrData();
const float *in0 = static_cast<const float *>(inputs_[0].Data().get());
const float *in1 = static_cast<const float *>(inputs_[1].Data().get());
float *out = static_cast<float *>(outputs_[0].MutableData());
auto num = outputs_[0].ElementNum();
for (int i = 0; i < num; ++i) {
out[i] = in0[i] + in1[i];
}
return RET_OK;
}
Custom算子属性解码样例
样例中是把属性里的字节流复制到了buf内。
auto prim = primitive_->value_as_Custom();
if (prim->attr()->size() < 1) {
return;
}
auto data_bytes = prim->attr()->Get(0)->data();
auto data_size = data_bytes->size();
char buf[100];
for (size_t i = 0; i < data_size; ++i) {
buf[i] = static_cast<char>(data_bytes->Get(i));
}
buf[data_size] = 0;
Custom算子注册
当前有提供的现成的宏REGISTER_CUSTOM_KERNEL可以进行算子注册，步骤如下：
TestCustomAddCreator函数用来创建Kernel。
通过宏REGISTER_CUSTOM_KERNEL进行算子注册，这里假定产商为BuiltInTest，算子类型为Add。
using mindspore::schema::PrimitiveType_AddFusion;
std::shared_ptr<Kernel> TestCustomAddCreator(const std::vector<tensor::MSTensor *> &inputs,
const std::vector<tensor::MSTensor *> &outputs,
const schema::Primitive *primitive, const lite::Context *ctx) {
return std::make_shared<TestCustomOp>(inputs, outputs, primitive, ctx);
}
constexpr auto kFloat32 = DataType::kNumberTypeFloat32;
REGISTER_CUSTOM_KERNEL(CPU, BuiltInTest, kFloat32, Add, TestCustomAddCreator)
Custom算子InferShape
整体实现与通用算子InferShape是一样的。步骤如下：
重载实现Infer函数，推导出output tensor的shape、format、data_type。
class TestCustomOpInfer : public KernelInterface {
public:
TestCustomOpInfer() = default;
~TestCustomOpInfer() = default;
Status Infer(std::vector<mindspore::MSTensor> *inputs, std::vector<mindspore::MSTensor> *outputs,
const schema::Primitive *primitive) override {
(*outputs)[0].SetFormat((*inputs)[0].format());
(*outputs)[0].SetDataType((*inputs)[0].DataType());
(*outputs)[0].SetShape((*inputs)[0].Shape());
return kSuccess;
}
};
Custom算子InferShape注册
当前有提供的现成的宏REGISTER_CUSTOM_KERNEL_INTERFACE可以进行Custom算子InferShape的注册，步骤如下：
CustomAddInferCreator函数用于创建自定义的KernelInterface。
通过宏REGISTER_CUSTOM_KERNEL_INTERFACE注册InferShape能力，这里的算子类型Add必须与REGISTER_CUSTOM_KERNEL时的算子类型一致。
std::shared_ptr<KernelInterface> CustomAddInferCreator() { return std::make_shared<TestCustomOpInfer>(); }
REGISTER_CUSTOM_KERNEL_INTERFACE(BuiltInTest, Add, CustomAddInferCreator)
自定义GPU算子
为支持GPU自定义算子的便捷开发，并使GPU自定义算子与内部的GPU算子共享一套资源，以加快调度效率，我们还提供了一套GPU相关的功能接口，相关API说明请参考mindspore::registry::opencl。 本文以样例代码解析的方式，向用户阐明自定义GPU算子开发的相关实现。用户需对如何实现自定义算子有所了解的情况下，再来阅读此文。 在代码仓样例代码中包含了对自定义GPU算子的实现、注册。
算子注册
本样例中注册的是Custom_Add自定义算子，关于该算子的创建与实现，请参考Custom算子定义和Custom算子实现。
实现创建算子实例的函数
实现自定义算子注册的第一步，需实现一个创建算子实例的函数。函数类型声明在include/registry/register_kernel.h，如下所示：
/// \brief CreateKernel Defined a functor to create a kernel.
///
/// \param[in] inputs Define input tensors of kernel.
/// \param[in] outputs Define output tensors of kernel.
/// \param[in] primitive Define attributes of op.
/// \param[in] ctx Define for holding environment variables during runtime.
///
/// \return Smart Pointer of kernel.
using CreateKernel = std::function<std::shared_ptr<kernel::Kernel>(
const std::vector<MSTensor> &inputs, const std::vector<MSTensor> &outputs, const schema::Primitive *primitive,
const mindspore::Context *ctx)>;
本例中实现的创建算子实例函数如下，函数返回一个CustomAddKernel类实例，该类为继承kernel::Kernel类的用户自定义算子类，关于该类的实现参考算子实现。
在函数内，除了将函数参数传递给CustomAddKernel类的构造函数外，还传递了一个布尔型的变量，该变量用于控制创建的CustomAddKernel实例处理的数据类型是FLOAT32还是FLOAT16。
namespace custom_gpu_demo {
std::shared_ptr<kernel::Kernel> CustomAddCreator(const std::vector<MSTensor> &inputs,
const std::vector<MSTensor> &outputs,
const schema::Primitive *primitive, const mindspore::Context *ctx) {
bool fp16_enable = false;
std::cout << "using fp32 add.\n" << std::endl;
return std::make_shared<CustomAddKernel>(inputs, outputs, primitive, ctx, fp16_enable);
}
}
注册算子
在注册GPU算子时，必须将设备类型声明为GPU，并将上一步实现的创建算子实例函数CustomAddCreator传入。
本样例注册了Custom_Add算子GPU内的Float32实现，注册代码如下所示，注册宏中的其他参数参考API说明。
const auto kFloat32 = DataType::kNumberTypeFloat32;
// Register custom "Custom_Add" operator
REGISTER_CUSTOM_KERNEL(GPU, BuiltInTest, kFloat32, Custom_Add, CustomAddCreator)
算子实现
在本样例中算子实现为CustomAddKernel类，该类继承mindspore::kernel::Kernel，重载实现必要的接口，从而实现自定义算子的运算。
构造及析构函数说明
在CustomAddKernel类构造函数中，保存了传递进来的布尔变量fp16_enable，并将其他参数传递给基类的构造函数。
在CustomAddKernel类析构函数中，调用FreeWeight()对因运算需要而临时申请的内存进行释放。
class CustomAddKernel : public kernel::Kernel {
public:
CustomAddKernel(const std::vector<MSTensor> &inputs, const std::vector<MSTensor> &outputs,
const schema::Primitive *primitive, const mindspore::Context *ctx,
bool fp16_enable)
: Kernel(inputs, outputs, primitive, ctx), fp16_enable_(fp16_enable) {}
~CustomAddKernel() override { FreeWeight(); }
...
}
类成员变量说明
opencl_runtime_
为OpenCLRuntimeWrapper类的实例，在算子内部可通过该对象调取MindSpore Lite提供的OpenCL操作相关接口mindspore::registry::opencl。
fp16_enable_
为算子是否使用FP16进行运算的标志。若要使用FP16进行运算，需将算子注册为FP16算子。本例中注册的是FP32算子。
weight_ptrs_
保存算子运算所需的临时内存的指针。
其他变量
其他变量为进行OpenCL操作时所需的变量，详细意义可查看OpenCL操作时对应的接口说明mindspore::registry::opencl。
class CustomAddKernel : public kernel::Kernel {
...
private:
const bool fp16_enable_;
cl::Kernel kernel_;
cl::Event event_;
cl::NDRange global_range_{cl::NullRange};
cl::NDRange local_range_{cl::NullRange};
std::vector<void *> weight_ptrs_;
registry::opencl::OpenCLRuntimeWrapper opencl_runtime_;
}
Prepare实现代码与说明
在图编译阶段mindspore::Model::Build，将调用算子的Prepare实现。用户可以在这里进行一些较为耗时，一次性的操作，以节约mindspore::Model::Predict时算子计算的时间。
在该样例中，通过重载Prepare接口，实现对自定义的OpenCL代码进行加载并编译。
检验环境
样例中，首先通过调用CheckSpecs，对算子的运行环境进行检查。
此处，在CheckSpecs中，检查了输入和输出的数据类型，及输入和输出的tensor数量。
通过MSTensor::IsConst()接口可以判断一个tensor的数据是否为常量，此处对非常量输入的数据类型，和算子注册时所声明处理的数据类型也进行了对比校验。对于常量数据的处理，参考本章后续的教程。
int Prepare() override {
auto ret = CheckSpecs();
if (ret != lite::RET_OK) {
std::cerr << "Prepare failed for check kernel specs!";
return ret;
}
...
}
int CheckSpecs() {
for (auto &tensor : inputs_) {
if (tensor.DataType() != DataType::kNumberTypeFloat32 && tensor.DataType() != DataType::kNumberTypeFloat16) {
std::cerr << "ArithmeticOpenCLKernel only support fp32/fp16 input";
return lite::RET_ERROR;
}
}
for (auto &tensor : outputs_) {
if (tensor.DataType() != DataType::kNumberTypeFloat32 && tensor.DataType() != DataType::kNumberTypeFloat16) {
std::cerr << "ArithmeticOpenCLKernel only support fp32/fp16 output";
return lite::RET_ERROR;
}
}
if (inputs_.size() != 2 || outputs_.size() != 1) {
std::cerr << "in size: " << inputs_.size() << ", out size: " << outputs_.size();
return lite::RET_ERROR;
}
for (int i = 0; i < inputs_.size(); ++i) {
auto &in_tensor = inputs_.at(i);
if (!in_tensor.IsConst()) {
if (fp16_enable_ && in_tensor.DataType() == mindspore::DataType::kNumberTypeFloat32) {
std::cerr << "Inputs data type error, expectation kNumberTypeFloat16 but kNumberTypeFloat32.";
return lite::RET_ERROR;
} else if (!fp16_enable_ && in_tensor.DataType() == mindspore::DataType::kNumberTypeFloat16) {
std::cerr << "Inputs data type error, expectation kNumberTypeFloat32 but kNumberTypeFloat16.";
return lite::RET_ERROR;
}
}
}
return lite::RET_OK;
}
加载自定义的OpenCL代码
通过opencl_runtime_调用OpenCLRuntimeWrapper::LoadSource接口加载自定义的OpenCL代码。
int Prepare() override {
...
const std::string kernel_name_ = "ElementAdd";
const std::string program_name = "Arithmetic";
std::string source = arithmetic_source;
if (opencl_runtime_.LoadSource(program_name, source) != kSuccess) {
std::cerr << "Load source failed.";
return lite::RET_ERROR;
}
...
}
arithmetic_source的为用户自定义的OpenCL代码，如下所示：
static const char *arithmetic_source =
"\n"
"#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n"
"__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;\n"
"\n"
"__kernel void ElementAdd(__read_only image2d_t input_a, __read_only image2d_t input_b, __write_only image2d_t "
"output,\n"
" const int2 output_shape) {\n"
" int X = get_global_id(0);\n"
" int Y = get_global_id(1);\n"
" if (X >= output_shape.x || Y >= output_shape.y) {\n"
" return;\n"
" }\n"
"\n"
" FLT4 a = READ_IMAGE(input_a, smp_none, (int2)(X, Y));\n"
" FLT4 b = READ_IMAGE(input_b, smp_none, (int2)(X, Y));\n"
" FLT4 result = a + b;\n"
"\n"
" WRITE_IMAGE(output, (int2)(X, Y), result);\n"
"}\n";
编译OpenCL代码
通过fp16_enable_指定不同的编译选项，以生成处理FLOAT16或FPLOAT32数据的代码。
通过opencl_runtime_调用OpenCLRuntimeWrapper::BuildKernel接口，得到编译后的cl::Kernel变量，保存在kernel_。
int Prepare() override {
...
std::vector<std::string> build_options_ext = {"-cl-mad-enable -cl-fast-relaxed-math -Werror"};
if (fp16_enable_) {
build_options_ext.push_back(" -DFLT4=half4 -DWRITE_IMAGE=write_imageh -DREAD_IMAGE=read_imageh");
} else {
build_options_ext.push_back(" -DFLT4=float4 -DWRITE_IMAGE=write_imagef -DREAD_IMAGE=read_imagef");
}
if (opencl_runtime_.BuildKernel(&kernel_, program_name, kernel_name_, build_options_ext) != kSuccess) {
std::cerr << "Build kernel failed.";
return lite::RET_ERROR;
}
...
}
设置OpenCL工作组和工作项
对注册为GPU的算子来说，除输入为常量的情况，所接收到的是Image格式的输入数据，Format为NHWC4（指C轴4字节对齐的NHWC格式数据）。
本例中也将所有数据转为这种格式进行计算和输出。
例程中实现的是一个简单的加法自定义算子，所以这里直接通过GpuTensorInfo函数计算输出数据Image内存所用宽度和高度来设置工作项。
int Prepare() override {
...
auto out_shape = GpuTensorInfo(&outputs_[0], &opencl_runtime_);
local_range_ = cl::NullRange;
global_range_ = cl::NDRange(out_shape.width, out_shape.height);
...
}
GpuTensorInfo的实现如下，首先通过Broadcast2GpuShape函数将tensor的shape转为四维，然后计算Format为NHWC4时的shape值。
再接着通过OpenCLRuntimeWrapper::GetMaxImage2DWidth及OpenCLRuntimeWrapper::GetMaxImage2DHeight接口得到Image内存所支持的最大宽度和高度，以此确定算子实际使用的Image内存宽度和高度。
struct GpuTensorInfo {
GpuTensorInfo() = default;
explicit GpuTensorInfo(const MSTensor *tensor, registry::opencl::OpenCLRuntimeWrapper *opencl_run) {
if (tensor == nullptr) {
return;
}
auto shape_ori = tensor->Shape();
int64_t shape[4];
Broadcast2GpuShape(shape, shape_ori.data(), shape_ori.size(), 1l);
N = shape[0];
H = shape[1];
W = shape[2];
C = shape[3];
Slice = UP_DIV(C, C4NUM);
if (tensor->DataType() == mindspore::DataType::kNumberTypeFloat16) {
FLT_size = sizeof(cl_half);
} else {
FLT_size = sizeof(cl_float);
}
FLT4_size = FLT_size * 4;
if (W * Slice <= opencl_run->GetMaxImage2DWidth()) {
height = N * H;
width = W * Slice;
} else {
height = N * H * W;
width = Slice;
if (height > opencl_run->GetMaxImage2DHeight()) {
height = -1;
width = -1;
}
}
ElementsNum = N * H * W * C;
Image2DSize = height * width * FLT4_size;
}
size_t N{1};
size_t H{1};
size_t W{1};
size_t C{1};
size_t Slice{};
size_t width{};
size_t height{};
size_t FLT_size{4};
size_t FLT4_size{16};
size_t ElementsNum{};
size_t Image2DSize{};
};
} // namespace
Broadcast2GpuShape的实现如下所示。
template <typename SrcT, typename DstT>
void Broadcast2GpuShape(DstT *dst, const SrcT *src, int src_num) {
if (src == nullptr || src_num <= 0) {
return;
}
auto *N = dst;
auto *H = dst + 1;
auto *W = dst + 2;
auto *C = dst + 3;
if (src_num == 1) { // 1 1 1 C
*C = src[0];
} else if (src_num == 2) { // N 1 1 C
*N = src[0];
*C = src[1];
} else if (src_num == 3) { // N 1 W C
*N = src[0];
*W = src[1];
*C = src[2];
} else if (src_num == 4) { // N H W C
*N = src[0];
*H = src[1];
*W = src[2];
*C = src[3];
} else if (src_num > 4) {
std::cerr << "GPU doesn't support ndim>=" << src_num;
}
}
template <typename SrcT, typename DstT>
void Broadcast2GpuShape(DstT *dst, const SrcT *src, int src_num, DstT default_value) {
for (int i = 0; i < 4; ++i) {
dst[i] = default_value;
}
if (src == nullptr || src_num <= 0) {
return;
}
Broadcast2GpuShape(dst, src, src_num);
}
将常量输入转为合适格式的数据，并分配GPU内存
对注册为GPU的算子来说，除输入为常量的情况，其它情况下，输入数据已经为Image格式的GPU内存数据。 为满足算子运算所需，用户需为常量输入设置合适的格式，必要时为其分配GPU内存。在此例，针对常量tensor的操作如下所示。
首先通过MSTensor::IsConst()接口判断输入是否为常量，并通过GpuTensorInfo计算转为Image格式时所需的内存大小。
然后分配该大小的局部内存weight，并通过PackNHWCToNHWC4函数将tensor内存转到weight中存储。
for (int i = 0; i < inputs_.size(); ++i) {
auto &in_tensor = inputs_.at(i);
if (in_tensor.IsConst()) {
GpuTensorInfo in_shape = GpuTensorInfo(&in_tensor, &opencl_runtime_);
std::vector<char> weight(in_shape.Image2DSize, 0);
bool src_is_fp16 = in_tensor.DataType() == mindspore::DataType::kNumberTypeFloat16;
PackNHWCToNHWC4(in_tensor.MutableData(), weight.data(), src_is_fp16, fp16_enable_, in_shape,
in_tensor.DataType());
...
PackNHWCToNHWC4函数实现如下，其中包含了对FLOAT16和FLOAT32类型的转换。
void PackNHWCToNHWC4(void *src, void *dst, bool src_is_fp16, bool dst_is_fp16, const GpuTensorInfo &tensor,
mindspore::DataType data_type) {
auto src_fp16 = reinterpret_cast<float16_t *>(src);
auto src_fp32 = reinterpret_cast<float32_t *>(src);
auto src_int32 = reinterpret_cast<int32_t *>(src);
auto dst_fp16 = reinterpret_cast<float16_t *>(dst);
auto dst_fp32 = reinterpret_cast<float32_t *>(dst);
auto dst_int32 = reinterpret_cast<int32_t *>(dst);
for (int n = 0, src_idx = 0; n < tensor.N; n++) {
for (int h = 0; h < tensor.H; ++h) {
for (int w = 0; w < tensor.W; ++w) {
for (int c = 0; c < tensor.C; ++c, ++src_idx) {
int dst_idx = ((n * tensor.H + h) * tensor.W + w) * tensor.Slice * C4NUM + c;
if (data_type == mindspore::DataType::kNumberTypeInt32) {
dst_int32[dst_idx] = src_int32[src_idx];
} else if (dst_is_fp16) {
dst_fp16[dst_idx] = src_is_fp16 ? src_fp16[src_idx] : static_cast<float16_t>(src_fp32[src_idx]);
} else {
dst_fp32[dst_idx] = src_is_fp16 ? static_cast<float32_t>(src_fp16[src_idx]) : src_fp32[src_idx];
}
}
}
}
}
if (tensor.ElementsNum == 1) {
if (dst_is_fp16) {
dst_fp16[3] = dst_fp16[2] = dst_fp16[1] = dst_fp16[0];
} else {
dst_fp32[3] = dst_fp32[2] = dst_fp32[1] = dst_fp32[0];
}
}
}
通过OpenCLRuntimeWrapper::GetAllocator得到分配GPU内存的内存分配器。
然后通过分配器的mindspore::Allocator::Malloc接口，可以申请到Image格式的GPU内存。
接着通过OpenCLRuntimeWrapper::WriteImage(void *buffer, void *src_data)接口，将已经转为NHWC4格式的weight数据写入到GPU内存中。
申请的GPU内存指针保存到weight_ptrs_中，以便在析构时释放。
DataType dtype =
fp16_enable_ ? mindspore::DataType::kNumberTypeFloat16 : mindspore::DataType::kNumberTypeFloat32;
auto allocator = opencl_runtime_.GetAllocator();
if (allocator == nullptr) {
std::cerr << "GetAllocator fail.";
FreeWeight();
return lite::RET_ERROR;
}
auto weight_ptr = allocator->Malloc(in_shape.width, in_shape.height, dtype);
if (weight_ptr == nullptr) {
std::cerr << "Malloc fail.";
FreeWeight();
return lite::RET_ERROR;
}
weight_ptrs_.push_back(weight_ptr);
if (opencl_runtime_.WriteImage(weight_ptr, weight.data()) != kSuccess) {
std::cerr << "WriteImage fail.";
FreeWeight();
return lite::RET_ERROR;
}
析构时调用的释放GPU内存函数如下，通过OpenCLRuntimeWrapper::GetAllocator得到分配GPU内存的内存分配器。
然后通过分配器的mindspore::Allocator::Free接口，可以释放申请到的GPU内存。
void FreeWeight() {
auto allocator = opencl_runtime_.GetAllocator();
if (allocator == nullptr) {
std::cerr << "GetAllocator fail.";
return;
}
for (auto &weight_ptr : weight_ptrs_) {
if (weight_ptr != nullptr) {
allocator->Free(weight_ptr);
weight_ptr = nullptr;
}
}
}
设置OpenCL内核运行时参数的值
某些OpenCL内核运行时不会改变的参数，可以在Prepare阶段进行设置。
在此例中，通过OpenCLRuntimeWrapper::SetKernelArg，设置ElementAdd运行时的第三个参数（计算的范围）。
int arg_idx = 3;
cl_int2 output_shape{static_cast<int>(global_range_[0]), static_cast<int>(global_range_[1])};
if (opencl_runtime_.SetKernelArg(kernel_, arg_idx, output_shape) != kSuccess) {
std::cerr << "Set kernel arg" << arg_idx << "failed.";
FreeWeight();
return lite::RET_ERROR;
}
ReSize及Execute实现代码与说明
通过重载实现Execute可以实现推理时算子的自定义运算操作。
调用
ReSize函数，以支持运行时shape变更
在本例中，首先调用PreProcess来处理运算前的一些准备工作。
在PreProcess()中，首先调用ReSize函数，该函数为需要用户重载实现的运行时shape变更适配接口。
在ReSize函数中，通过调用CheckOutputs判断算子的输出tensor的shape是否存在非法值，以判断是否需要重新进行shape推理。若不需要，直接返回。
在需要进行shape推理时，通过registry::RegisterKernelInterface::GetKernelInterface获得该算子所注册的shape推理函数，此处得到的其实就是本例程中用户实现并注册的InferShape函数。
在重新推理之后，通过调用之前实现的Prepare接口，重新申请和分配算子运算时需要的内存及相关变量。
int ReSize() override {
if (CheckOutputs(outputs_) == lite::RET_OK) {
return lite::RET_OK;
}
auto status =
registry::RegisterKernelInterface::GetKernelInterface("", primitive_)->Infer(&inputs_, &outputs_, primitive_);
if (status != kSuccess) {
std::cerr << "infer failed." << std::endl;
return lite::RET_ERROR;
}
ret = Prepare();
if (ret != lite::RET_OK) {
std::cerr << "ReSize failed for kernel prepare!";
return ret;
}
return lite::RET_OK;
}
int PreProcess() {
int ret;
ret = ReSize();
if (ret != lite::RET_OK) {
return ret;
}
...
}
int Execute() override {
if (inputs_.size() != 2) {
return lite::RET_PARAM_INVALID;
}
PreProcess();
...
}
为输出tensor申请内存分配
在算子运行前，需要为输出tensor申请分配GPU内存，由于框架的限制，该GPU内存需要托管给框架管理，用户不可人为释放。具体操作流程如下：
通过调用输出tensor的
allocator()接口，可以得到框架中管理这个tensor的内存分配器，在GPU注册算子中，则为负责分配GPU内存的内存分配器。计算需要分配的内存大小，此例中通过
GpuTensorInfo函数来计算。通过内存分配器的
Malloc接口申请内存，用户可分别通过void *Malloc(size_t weight, size_t height, DataType type)和void *Malloc(size_t size)接口得到Image或Buffer格式的内存。通过
SetData接口，将申请的内存赋值给tensor，此后，此内存将由框架统一管理，用户不可手动释放。
int PreProcess() {
...
for (auto i = 0; i < outputs_.size(); ++i) {
auto *output = &outputs_.at(i);
auto img_info = GpuTensorInfo(output, &opencl_runtime_);
auto allocator = output->allocator();
if (allocator == nullptr) {
std::cerr << "The output tensor of OpenCL kernel must have an allocator.";
return lite::RET_ERROR;
}
auto data_ptr = allocator->Malloc(img_info.width, img_info.height, output->DataType());
if (data_ptr == nullptr) {
std::cerr << "Malloc data failed";
return lite::RET_ERROR;
}
output->SetData(data_ptr);
}
return lite::RET_OK;
}
运行OpenCL内核
通过SetKernelArg接口设置OpenCL的Kernel运行时的参数，通过RunKernel运行OpenCL的Kernel。
int Execute() override {
...
std::cout << this->name() << " Running!" << std::endl;
auto input_0_ptr = weight_ptrs_[0] == nullptr ? inputs_[0].MutableData() : weight_ptrs_[0];
auto input_1_ptr = weight_ptrs_[1] == nullptr ? inputs_[1].MutableData() : weight_ptrs_[1];
int arg_idx = 0;
if (opencl_runtime_->SetKernelArg(kernel_, arg_idx++, input_0_ptr) != kSuccess) {
std::cerr << "Set kernel arg" << arg_idx - 1 << "failed.";
return lite::RET_ERROR;
}
if (opencl_runtime_->SetKernelArg(kernel_, arg_idx++, input_1_ptr) != kSuccess) {
std::cerr << "Set kernel arg" << arg_idx - 1 << "failed.";
return lite::RET_ERROR;
}
if (opencl_runtime_->SetKernelArg(kernel_, arg_idx++, outputs_[0].MutableData()) != kSuccess) {
std::cerr << "Set kernel arg" << arg_idx - 1 << "failed.";
return lite::RET_ERROR;
}
if (opencl_runtime_->RunKernel(kernel_, global_range_, local_range_, nullptr, &event_) != kSuccess) {
std::cerr << "Run kernel failed.";
return lite::RET_ERROR;
}
return lite::RET_OK;
}