如何給MindSpore添加一個新的硬件后端？快速構建測試環境！

摘要：介紹如何給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 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_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";  // 添加新的硬件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 添加文件夾）:

下面介紹要創建的針對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的compile graph過程中）運行一次， 而LaunchKernel(..)會在每次圖執行的過程中被調用。例如跑一個CNN的推理, 需要infernce 64張圖片，網絡的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 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的深度學習！

 

點擊關注，第一時間了解華為雲新鮮技術~