写点什么

如何给 MindSpore 添加一个新的硬件后端?快速构建测试环境!

发布于: 4 小时前

摘要:介绍如何给 MindSpore 添加一个新的硬件后端。


本文分享自华为云社区《如何给MindSpore添加一个新的硬件后端?快速构建测试环境!》,原文作者:HWCloudAI。


MindSpore 是华为自研的新一代 AI 开源计算框架。最佳匹配昇腾 AI 处理器算力的全场景深度学习框架,为数据科学家和算法工程师提供设计友好运行高效的开发体验推动人工智能软硬件应用生态繁荣发展


MindSpore 支持异构算力,除支持华为自研的达芬奇架构的 Ascend NPU 外还支持 CPU(e.g. MKLDNN) 以及 GPU(e.g. CUDA kernels)算子的运行。(注:MindSpore 支持整网在不同的硬件平台上运行,并不支持同一张网络的不同 partition 在不同的硬件平台上运行,这点和 TensorFlow 的 graph partition 异构运行模式不一样)。


当前 AI 芯片行业“热闹非凡”,国内外,大小新老厂商都在推出自己的 AI 加速芯片。现在大家都应该看得很清楚,硬件要想成功,离不开软件栈及生态的支撑。MindSpore 不仅为支撑华为的 AI 软硬件栈服务,更想在整个 AI 生态中占据自己的一片天地。


MindSpore 目前还处于推广和发展完善阶段,本文想抛砖引玉介绍如何给 MindSpore 添加一个新的硬件后端,同时对 MindSpore 源代码的目录结构也做一些基本介绍,希望能为国内外的 AI 硬件厂商和感兴趣的开发人员提供一些有用信息和参考,让大家能来共同使用 MindSpore 作为测试和对接 AI 芯片的框架快速构建整网模型的测试环境


本文针对的是 MindSpore r1.1 版本的源代码:https://gitee.com/mindspore/mindspore/tree/r1.1/对于如何从源码编译及安装 MindSpore,以及对于相关软件版本的需求,请参考:https://www.mindspore.cn/install/

测试用例 


本文将针对一个简单的 Dense layer 网络:https://www.mindspore.cn/doc/api_python/zh-CN/r1.1/mindspore/nn/mindspore.nn.Dense.html#mindspore.nn.Dense来示范如何让这个 layer 运行在一个新的硬件后端上。注:本文针对的是基本的静态图执行模式:https://www.mindspore.cn/doc/programming_guide/zh-CN/r1.1/context.html


import mindsporeimport numpy as npimport mindspore.nn as nnfrom mindspore import context, Tensor
context.set_context(device_target="CPU", mode=context.GRAPH_MODE)
# 32, 16net = nn.Dense(32, 16, weight_init='ones', bias_init=1.2)#, activation='relu')
# 48, 32input_data = Tensor(np.ones([48, 32]).astype(np.float32), mindspore.float32)output = net(input_data)
print(output.asnumpy())
复制代码


注:在这里我注释掉了 activation 的 ReLU,所以此 Dense layer 就相当于一个只有 2 个 node 的小网络(MatMul +BiasAdd) 此用例的结果是一个 48 * 16 的二维矩阵,每个 element 的值都是 33.2)

此文将以从上到下的流程,介绍 MindSpore 支持一个新硬件后端所需要修改的组件。我们这里将需要支持的新硬件称为 XPU, 我们想要达到的修改 MindSpore 代码后的效果是将上述用例中的 device_target 改为 XPU, 并在让 Dense layer 在加速器 XPU 上运行。e.g.


context.set_context(device_target="XPU", mode=context.GRAPH_MODE)
复制代码


注:此文不会展示具体类和函数的实现细节,具体的实现可以参考相对应目录下已支持的硬件后端的实现,例如:CPU, GPU, Ascend

添加新的 device target 参数选项


首先从前端 MEPython 层需要添加新的 valid_targets:https://gitee.com/mindspore/mindspore/blob/r1.1/mindspore/context.py


def set_device_target(self, target):        valid_targets = ["CPU", "GPU", "Ascend", "Davinci", "XPU"] # 将新的后端添加到此list中        if not target in valid_targets:            raise ValueError(f"Target device name {target} is invalid! It must be one of {valid_targets}")        if target == "Davinci":            target = "Ascend"        self.set_param(ms_ctx_param.device_target, target)        if self.enable_debug_runtime and target == "CPU":            self.set_backend_policy("vm") 
复制代码


接着需要在 C++的 ms context 组件中添加新的 target:https://gitee.com/mindspore/mindspore/blob/r1.1/mindspore/core/utils/ms_context.h


const int kGraphMode = 0;const int kPynativeMode = 1;const char kCPUDevice[] = "CPU";const char kGPUDevice[] = "GPU";const char kXPUDevice[] = "XPU";  // 添加新的硬件targetconst char kAscendDevice[] = "Ascend";const char kDavinciInferenceDevice[] = "AscendInference";const char kDavinciDevice[] = "Davinci";const char KNpuLog[] = "_npu_log";const unsigned int MAX_CALL_DEPTH_DEFAULT = 1000;
// 添加新的硬件到以下set中const std::set<std::string> kTargetSet = {kCPUDevice, kGPUDevice, kXPUDevice, kAscendDevice, kDavinciDevice};
复制代码


添加新的 runtime device


在 runtime device 目录下:https://gitee.com/mindspore/mindspore/tree/r1.1/mindspore/ccsrc/runtime/device是和各个具体后端硬件特性相关的组件,例如:device 端的地址空间,device 端的内存管理(分配,回收),kernel runtime 组件等, 还有硬件 device 相关的一些通讯组件,例如支持分布式通讯的 MPI 组件。我们首先在下图中的目录下添加一个叫 xpu 的文件夹 (注意修改 CMakeLists.txt 添加文件夹):



下面介绍要创建的针对 xpu 加速器 3 个新的基本组件:


xpu_device_address :主要表示加速器 device 侧的内存地址信息,以及 host 端和 device 端之间内存搬移的 API 接口,例如在 NVIDIA GPU 上可以是 wrapper of:cudaMemcpyAsyncxpu_device_address.h


#include <string>#include <vector>#include "runtime/device/device_address.h"#include "utils/shape_utils.h"
namespace mindspore {namespace device {namespace xpu {class XPUDeviceAddress : public DeviceAddress { public: XPUDeviceAddress(void *ptr, size_t size) : DeviceAddress(ptr, size) {}
XPUDeviceAddress(void *ptr, size_t size, const string &format, TypeId type_id) : DeviceAddress(ptr, size, format, type_id) {}
~XPUDeviceAddress() override = default;
bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override; bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override; DeviceAddressType DeviceType() const override { return DeviceAddressType::kXPU; }};} // namespace xpu} // namespace device} // namespace mindspore
复制代码


xpu_resource_manager: 主要负责 device 端的内存和其他资源的管理,分配和调度。xpu_resource_manager.h


#include <vector>#include <map>#include "backend/session/kernel_graph.h"#include "backend/session/session_basic.h"#include "runtime/device/device_address.h"#include "runtime/device/xpu/xpu_simple_mem_plan.h"namespace mindspore {namespace device {namespace xpu {class XPUResourceManager { public:  XPUResourceManager() = default;  ~XPUResourceManager();
void AssignMemory(const session::KernelGraph *graph); void IncreaseAddressRefCount(const session::KernelGraph *graph); void DecreaseAddressRefCount(const AnfNodePtr &kernel); void *MemMalloc(size_t mem_size); void MemFree(void *ptr);
private: void MemFree(); XPUSimpleMemPlan mem_plan_;
size_t mem_size_{0}; uint8_t *mem_ptr_{nullptr}; bool dynamic_malloc_{false}; std::map<void *, size_t> dynamic_mem_;};} // namespace xpu} // namespace device} // namespace mindspore
复制代码


xpu_kernel_runtime:硬件算子的执行控制模块,主要负责硬件 runtime 的启动(Init()),网络在硬件上的执行(Run(..)),已经硬件执行完后的清理工作(ReleaseDeviceRes())xpu_kernel_runtime.h


#include <memory>#include <vector>#include <string>#include <map>#include <set>#include "runtime/device/kernel_runtime.h"#include "runtime/device/kernel_runtime_manager.h"#include "backend/session/kernel_graph.h"#include "backend/session/session_basic.h"#include "runtime/device/xpu/xpu_resource_manager.h"#include "backend/session/anf_runtime_algorithm.h"#include "utils/any.h"namespace mindspore {namespace device {namespace xpu {class XPUKernelRuntime : public KernelRuntime { public:  XPUKernelRuntime() = default;  ~XPUKernelRuntime() override = default;
bool Init() override; void ReleaseDeviceRes() override; bool Run(session::KernelGraph *graph, bool is_task_sink) override; void AssignKernelAddress(session::KernelGraph *kernel_graph); void CreateOutputTensors(session::KernelGraph *kernel_graph, const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs); void BindInputOutput(session::KernelGraph *kernel_graph, const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs);
protected: bool SyncStream() override { return true; }; DeviceAddressPtr CreateDeviceAddress(void *device_ptr, size_t device_size, const string &format, TypeId type_id) override;
private: XPUResourceManager resource_manager_; std::set<DeviceAddressPtr> bound_addresses_; std::map<AnfNodePtr, tensor::TensorPtr> input_param_tensor_map_;};
MS_REG_KERNEL_RUNTIME(kXPUDevice, XPUKernelRuntime);
} // namespace xpu} // namespace device} // namespace mindspore
复制代码


添加新的 target session


MindSpore 的 Session(会话)提供了 Op kernel 执行和 Tensor 求值的环境。Session 是控制代表神经网络的数据流图的核心模块。它主要有图编译(kernel 生成),图优化,和图执行三个主要步骤。MindSpore 针对每个后端硬件平台都会有自己的 Session 组件,相关代码在 backend/session 这个目录中:https://gitee.com/mindspore/mindspore/tree/r1.1/mindspore/ccsrc/backend/session

我们针对 xpu 创建新的 session 类:xpu_session.h


#include <string>#include <memory>#include <map>#include <vector>#include "backend/session/session_basic.h"#include "backend/session/kernel_graph.h"#include "runtime/device/xpu/xpu_kernel_runtime.h" // use the new xpu kernel runtime#include "backend/session/session_factory.h"namespace mindspore {namespace session {class XPUSession : public SessionBasic { public:  XPUSession() = default;  ~XPUSession() override = default;  void Init(uint32_t device_id) override { InitExecutor(kXPUDevice, device_id); }
GraphId CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) override; void RunGraphImpl(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs) override; void Optimize(const std::shared_ptr<KernelGraph> &kernel_graph);
protected: void UnifyMindIR(const KernelGraphPtr &graph) override { return; } void CreateOutputTensors(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &input_tensors, VectorRef *, std::map<tensor::TensorPtr, session::KernelWithIndex> *tensor_to_node) override;
private: void SetKernelInfo(const KernelGraph *kernel_graph); void BuildKernel(const KernelGraph *kernel_graph); device::xpu::XPUKernelRuntime *runtime_ = dynamic_cast<device::xpu::XPUKernelRuntime*>(device::KernelRuntimeManager::Instance().GetKernelRuntime(kXPUDevice, 0));};MS_REG_SESSION(kXPUDevice, XPUSession);} // namespace session} // namespace mindspore
复制代码


在图编译(CompileGraphImpl(..))的步骤中,主要是要生成(BuildKernel(..))表示神经网络数据流图中的每个节点 op 相对应的 kernel,并保存每个节点的 kernel 信息在图中(SetKernelInfo(..)),以供在后面的图执行(RunGraphImpl(..))步骤中被调用。

添加针对新硬件的 kernel


MindSpore 所支持的硬件后端对于各个 op 算子的支持在 backend/kernel_compiler 目录下:https://gitee.com/mindspore/mindspore/tree/r1.1/mindspore/ccsrc/backend/kernel_compiler



在这里我们可以看到针对不多的硬件后端,每一个文件夹代表着不同 kernel 的类型,其中:


cpu:里面有调用 MKLDNN(oneDNN) 的算子,也有纯 c++写的算子。


gpu: 里面有调用 cudnn/cublas 的算子,也有用 cuda 写的算子,还有支持分布式训练与 NCCL 相关的算子。


Ascend: 与华为达芬奇 AI 芯片相关的算子 kernel 文件夹有:tbe, aicpu,akg,hccl 等


下面来介绍为我们的新硬件后端添加 kernel 支持所需的组件,首先在上面的目录下创建一个叫 xpu 的文件夹 (注意修改 CMakeLists.txt 添加文件夹)在新文件夹中我们首先来创建针对 xpu kernel 的基类:xpu_kernel.h:


#include <string>#include <vector>#include <memory>#include <numeric>#include <functional>#include "backend/kernel_compiler/kernel.h"#include "ir/anf.h"#include "backend/session/anf_runtime_algorithm.h"#include "utils/ms_utils.h"
using mindspore::kernel::Address;using mindspore::kernel::AddressPtr;namespace mindspore {namespace kernel {
class XPUKernel : public kernel::KernelMod { public: XPUKernel() = default; ~XPUKernel() override = default;
void Init(const CNodePtr &kernel_node); virtual void InitKernel(const CNodePtr &kernel_node) = 0; bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace, const std::vector<AddressPtr> &outputs, void * stream_ptr) override { return Launch(inputs, workspace, outputs); };
virtual bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace, const std::vector<AddressPtr> &outputs) = 0; const std::vector<size_t> &GetInputSizeList() const override { return input_size_list_; } const std::vector<size_t> &GetOutputSizeList() const override { return output_size_list_; } const std::vector<size_t> &GetWorkspaceSizeList() const override { return workspace_size_list_; }
void SetOpName(const std::string &op_name) { op_name_ = op_name; } const std::string GetOpName() const { return op_name_; }
protected: virtual void InitInputOutputSize(const CNodePtr &kernel_node); std::vector<size_t> input_size_list_ = {}; std::vector<size_t> output_size_list_ = {}; std::vector<size_t> workspace_size_list_ = {};
std::string bin_path_ = {}; std::string tilingName_ = {};
};} // namespace kernel} // namespace mindspore
复制代码


现在流行的框架对于算子 kernel 的支持普遍是采用以算子名(opcode)来命名 kernel,例如 mindspore 里 mkldnn 的 cpu kernels:MindSpore/mindspore 这种形式的优点是 repo 代码文件很清晰,每个算子的特定属性可以很方便的表达。缺点是会有可能造成一些 duplicate 的代码逻辑。由于本文针对的用例很简单,实际上只需要支持 2 个算子:MatMul 和 BiasAdd,我们将采用按输入输出 Tensor 个数来命名的 kernel 类实现方式。


由于 MatMul 和 BiasAdd 都是 2 个输入 1 个输出的算子,我们定义我们的 kernel 类名为:two_in_one_out_xpu_kernel.h


#include "backend/kernel_compiler/xpu/xpu_kernel.h" // xpu kernel base class#include "backend/kernel_compiler/xpu/xpu_kernel_factory.h"
#include <stdio.h>#include <limits.h>#include <stdlib.h>#include <unistd.h>#include <fcntl.h>#include <dirent.h>#include <algorithm>
#include <fstream>#include <iostream>
namespace mindspore {namespace kernel {
class TwoInOneOutXPUKernel : public XPUKernel { public: TwoInOneOutXPUKernel() = default; ~TwoInOneOutXPUKernel() override = default;
void InitKernel(const CNodePtr &kernel_node) override;
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace, const std::vector<AddressPtr> &outputs) override;
private: bool NeedsFormatTransformation();
char trans_a_{TRANSPOSE_NO}; char trans_b_{TRANSPOSE_NO}; int32_t dim_m_{0}; int32_t dim_n_{0}; int32_t dim_k_{0};
std::vector<size_t> inputA_shape_; std::vector<size_t> inputB_shape_; std::vector<size_t> output_shape_;
size_t input_a_size_ = 0; size_t input_b_size_ = 0; size_t output_size_ = 0;
void *inputA_data_ = nullptr; void *inputB_data_ = nullptr; void *output_data_ = nullptr;};
MS_REG_XPU_KERNEL( TwoInOneOutXPU, mindspore::device::xpu::KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), TwoInOneOutXPUKernel);} // namespace kernel} // namespace mindspore
复制代码


在这里我们有使用到"backend/kernel_compiler/xpu/xpu_kernel_factory.h" 对于 kernel 工厂类的创建我们就不细述,具体可以参考 cpu_kernel_factory.h:https://gitee.com/mindspore/mindspore/blob/r1.1/mindspore/ccsrc/backend/kernel_compiler/cpu/cpu_kernel_factory.h


对于每个 kernel 最基本的 2 个 function 就是 InitKernel(..)和 LaunchKernel(..) 分别负责 kernel 的初始化和运行。这里需要注意的是,对于一般像 CNN 静态图的执行,InitKernel(..)只会在 kernel 创建时(上述 session 的 compilegraph 过程中)运行一次, 而 LaunchKernel(..)会在每次图执行的过程中被调用。例如跑一个 CNN 的推理, 需要 infernce64 张图片,网络的 batch size is 32, 那整张图需要被执行 2 遍,也就是说针对每个 kernel,InitKernel(..)会被调用 1 次,而 LaunchKernel(..)会被调用 2 次。


我们这里不细述 MatMul 和 BiasAdd kernel 的具体实现,只介绍一些 MindSpore 里针对算子 kernel 所需要使用的一些基本 API:


获取 TwoInOneOutXPUKernel 的 input,output shape 信息:


inputA_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 0);inputB_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 1);output_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, 0);
复制代码


获取算子属性信息,e.g. MatMul 的转置信息:


bool trans_a = AnfAlgo::GetNodeAttr<bool>(kernel_node, TRANSPOSE_A);bool trans_b = AnfAlgo::GetNodeAttr<bool>(kernel_node, TRANSPOSE_B);
复制代码


在 Launch 里获得输入,输出 memory 的指针:


auto input_a = reinterpret_cast<float *>(inputs[0]->addr);auto input_b = reinterpret_cast<float *>(inputs[1]->addr);auto output = reinterpret_cast<float *>(outputs[0]->addr);
复制代码


其他注意事项


和其他主流框架一样,MindSpore 里也会有一些自己的标准和规范,下面介绍一些自己踩过的“坑”和大家分享:


MindSpore 里的 Tensor 的默认 format 是 NCHW。如果你所添加的硬件后端所支持的格式不一样,要注意添加格式转换。格式转换可以在每个 kernel 的调用前后去做(效率差), 也可以利用图优化 pass, 以整个网络为视野来高效的插入格式转换节点。


精度转换,如果你的硬件平台只支持某些精度,例如 fp16,而网络是 fp32 那就要注意精度的转换,精度转换和上述格式转换类似。精度转换可以在 host 端做,也可以在 device 端做(如果硬件支持)。


对于每个 kernel 的代码逻辑要区别哪些 data 是不变的,哪些是会变的,需要每次执行前重新初始化的,这样可以合理和正确的分配不同逻辑代码去相应 InitKernel(..) 或 LaunchKernel(..)里去。


对于某些 Python 前端的 LayerAPI,MindSpore 有自己的一些属性设置,例如对于 Denselayer:https://gitee.com/mindspore/mindspore/blob/r1.1/mindspore/nn/layer/basic.py的第 2 个输入矩阵是被转置过的:


self.matmul = P.MatMul(transpose_b=True)self.batch_matmul = P.BatchMatMul(transpose_b=True)self.activation = get_activation(activation) if isinstance(activation, str) else activationif activation is not None and not isinstance(self.activation, (Cell, Primitive)):    raise TypeError("The activation must be str or Cell or Primitive,"" but got {}.".format(activation))self.activation_flag = self.activation is not None
复制代码


对于 Debug,可以添加下面的环境变量来帮助输出信息:


export GLOG_v=1export SLOG_PRINT_TO_STDOUT=1
复制代码


对于 CMake 文件的修改,可以在开始测试时把新添加的文件都添加在 if (ENABLE_CPU)下,CPU 对于 MindSpore 相当于一个基线平台,也就是说无论是你 build GPU 还是华为的 D/Ascend target, CPU 相关的文件都会被 build。

总结


本文是作者根据自己对于 MindSpore 的理解,和大家分享的一个如何修改 MindSpore 源码来添加一个新硬件后端的技术文章。一个开源软件框架的成功,离不开社区的支持和各个厂商的参与,希望本文能启到一个抛砖引玉的作用,让更多的硬件厂商和开发者也能参与到 MindSpore 的生态发展中来。也欢迎大家拍砖来一起讨论!


了解完 MindSpore 的关键技术是不是很心动呢!赶紧【点击链接】并【立即报名】,即可在 ModelArts 平台学习到一个经典案例掌握基于 MindSpore 的深度学习!


点击关注,第一时间了解华为云新鲜技术~

发布于: 4 小时前阅读数: 3
用户头像

提供全面深入的云计算技术干货 2020.07.14 加入

华为云开发者社区,提供全面深入的云计算前景分析、丰富的技术干货、程序样例,分享华为云前沿资讯动态,方便开发者快速成长与发展,欢迎提问、互动,多方位了解云计算! 传送门:https://bbs.huaweicloud.com/

评论

发布
暂无评论
如何给MindSpore添加一个新的硬件后端?快速构建测试环境!