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

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 mindspore
import numpy as np
import mindspore.nn as nn
from mindspore import context, Tensor

context.set_context(device_target="CPU", mode=context.GRAPH_MODE)

# 32, 16
net = nn.Dense(32, 16, weight_init='ones', bias_init=1.2)#, activation='relu')

# 48, 32
input_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参数选项

首先从前端ME Python层需要添加新的valid_targetshttps://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";  // 添加新的硬件target
const 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 添加文件夹):

1.png

下面介绍要创建的针对xpu加速器3个新的基本组件:
· xpu_device_address :主要表示加速器device侧的内存地址信息,以及host端和device端之间内存搬移的API接口,例如在NVIDIA GPU上可以是wrapper of:cudaMemcpyAsync
xpu_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

2.png

在这里我们可以看到针对不多的硬件后端,每一个文件夹代表着不同kernel的类型,其中:
· cpu:里面有调用MKLDNN(oneDNN) 的算子,也有纯c++写的算子。
· gpu: 里面有调用cudnn/cublas的算子,也有用cuda写的算子,还有支持分布式训练与NCCL相关的算子。
· Ascend: 与华为达芬奇AI芯片相关的算子kernel文件夹有:tbeaicpu,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个算子:MatMulBiasAdd,我们将采用按输入输出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的compile graph过程中)运行一次, 而LaunchKernel(..)会在每次图执行的过程中被调用。例如跑一个CNN的推理, 需要infernce 64张图片,网络的batch size is 32, 那整张图需要被执行2遍,也就是说针对每个kernel,InitKernel(..)会被调用1次,而LaunchKernel(..)会被调用2次。

我们这里不细述MatMulBiasAdd 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 activation
if 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=1
export SLOG_PRINT_TO_STDOUT=1

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


总结

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


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







(完)