CANN/Ascend Boost Comm自定义算子开发示例

Ascend Boost Comm 库开发单算子示例【免费下载链接】ascend-boost-comm算子公共平台南向对接不同组织开发的算子库北向支撑不同加速库应用实现M x N算子能力复用项目地址: https://gitcode.com/cann/ascend-boost-comm本教程以add算子为例提供能够基于Ascend Boost Comm运行的算子编写教程.算子功能两个输入张量在指定维度相加成一个输出张量。新增文件在examples/ops下新增addcustom目录该目录下主要存放算子实现部分的代码文件具体内容见后文目录结构如下addcustom ├── op_kernel // device侧实现文件(包括核函数入口、以及实现文件) │ └── addcustom.cpp ├── tiling // 新增算子tiling │ ├── addcustom_tiling.cpp // tiling实现核心算法 │ ├── addcustom_tiling.h // 算子tiling接口 │ └── tiling_data.h // tiling和kernel传递结构体tiling_data的定义 ├── CMakeLists.txt // 新增算子编译CMake ├── addcustom_kernel.cpp // 校验 └── addcustom_operation.cpp // shape验证新增example/include/asdops/params/addcustom.h文件定义Addcustom操作的参数结构体内容如下#ifndef ATBOPS_PARAMS_ADDCUSTOM_H #define ATBOPS_PARAMS_ADDCUSTOM_H #include cstdint #include string #include sstream #include mki/utils/SVector/SVector.h namespace Mki { namespace OpParam { struct Addcustom { bool operator(const Addcustom other) const { (void)other; return true; } }; } // namespace OpParam } // namespace Mki #endif修改文件在example/include/asdops/params/params.h中包含新增的头文件example/include/asdops/params/addcustom.h:#include atbops/params/addcustom.h环境准备您可参考环境准备进行编译编译和测试环境搭建环境准备好之后就可以开始算子开发之旅。算子实现算子实现主要包括device侧算子实现和host侧tiling实现。tiling开发tiling开发的核心概念TilingData、Workspace、TilingKey、BlockDim等可访问术语表-昇腾社区查看。tiling_data.h文件路径example/ops/addcustom/tiling/tiling_data.h主要功能描述执行tiling计算所需的结构化数据定义。#ifndef ASCEND_OPS_ADDCUSTOM_TILING_DATA #define ASCEND_OPS_ADDCUSTOM_TILING_DATA #include cstdint namespace Mki { struct AddcustomTilingData { uint32_t totalLength; // 总数据长度 uint32_t tileNum; // Tiling 块数 }; } #endifaddcustom_tiling.h文件路径example/ops/addcustom/tiling/addcustom_tiling.h主要功能tiling过程主要是完成数据的切分这里则是函数声明。#ifndef ASCEND_OPS_ADDCUSTOM_TILING_H #define ASCEND_OPS_ADDCUSTOM_TILING_H #include mki/launch_param.h #include mki/kernel_info.h #include mki/utils/status/status.h namespace Mki { Status AddcustomTiling(const LaunchParam launchParam, KernelInfo kernelInfo); } #endifaddcustom_tiling.cpp文件路径example/ops/addcustom/tiling/addcustom_tiling.cpp主要功能实现切分功能的主体函数(具体函数实现请参考上面文件路径下的内容)#include addcustom_tiling.h #include mki/utils/assert/assert.h #include mki/utils/log/log.h #include mki/utils/platform/platform_info.h #include mki/utils/math/math.h #include mki/utils/SVector/SVector.h #include atbops/params/addcustom.h #include tiling_data.h // 定义最小的块长度 constexpr uint32_t MIN_BLOCK_LENGTH 32; namespace Mki { Status AddcustomTiling(const LaunchParam launchParam, KernelInfo kernelInfo) { // 具体函数实现请参考上面文件路径下的内容 } } // namespace Mkikernel开发kernel相关的Compute、CopyIn、CopyOut等概念可访问术语表-昇腾社区查看。addcustom.cpp文件路径example/ops/addcustom/op_kernel/addcustom.cpp主要功能根据tiling信息完成所有数据的搬运和计算。#include kernel_operator.h #include ops/utils/common/kernel/kernel_utils.h #include ops/addcustom/tiling/tiling_data.h static constexpr uint32_t BUFFER_NUM 1; static constexpr uint32_t MAX_UB_SIZE 188 * 1024; // double buffer, 每块94KB共188KB class Addcustom { public: __aicore__ inline Addcustom() {} __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength, uint32_t tileNum) { // 具体函数实现请参考上面文件路径下的内容 } __aicore__ inline void Process() { // 具体函数实现请参考上面文件路径下的内容 } private: __aicore__ inline void CopyIn(int32_t progress) { // 具体函数实现请参考上面文件路径下的内容 } __aicore__ inline void Compute(int32_t progress) { // 具体函数实现请参考上面文件路径下的内容 } __aicore__ inline void CopyOut(int32_t progress) { // 具体函数实现请参考上面文件路径下的内容 } private: AscendC::TPipe pipe; AscendC::TQueAscendC::QuePosition::VECIN, BUFFER_NUM inQueueX, inQueueY; AscendC::TQueAscendC::QuePosition::VECOUT, BUFFER_NUM outQueueZ; AscendC::GlobalTensorhalf xGm; AscendC::GlobalTensorhalf yGm; AscendC::GlobalTensorhalf zGm; uint32_t blockLength; uint32_t tileNum; uint32_t tileLength; }; inline __aicore__ void InitTilingData(const __gm__ uint8_t *p_tilingdata, Mki::AddcustomTilingData *tilingdata) { // 具体函数实现请参考上面文件路径下的内容 } extern C __global__ __aicore__ void addcustom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR tiling) { // 具体函数实现请参考上面文件路径下的内容 }addcustom_kernel.cpp文件路径example/ops/addcustom/addcustom_kernel.cpp主要功能启动device侧实现前对输入和输出进行了检查以及进行device侧的初始化。#include mki/base/kernel_base.h #include mki_loader/op_register.h #include mki/utils/assert/assert.h #include mki/utils/log/log.h #include atbops/params/params.h #include ops/addcustom/tiling/addcustom_tiling.h #include ops/addcustom/tiling/tiling_data.h namespace Mki { class AddcustomKernel : public KernelBase { public: explicit AddcustomKernel(const std::string kernelName, const BinHandle *handle) noexcept : KernelBase(kernelName, handle) { } /* --------- 框架回调 --------- */ bool CanSupport(const LaunchParam launchParam) const override { // 具体函数实现请参考上面文件路径下的内容 } uint64_t GetTilingSize(const LaunchParam launchParam) const override { // 具体函数实现请参考上面文件路径下的内容 } Status InitImpl(const LaunchParam launchParam) override { // 具体函数实现请参考上面文件路径下的内容 } }; /* ---------- 注册到框架 ---------- */ REG_KERNEL_BASE(AddcustomKernel); } // namespace Mkiaddcustom_operation.cpp文件路径example/ops/addcustom/addcustom_operation.cpp主要功能定义 Operation 行为Operation 是 Ascend Boost Comm 框架下 op 的最高程度抽象包括选择最佳 kernel 的策略。#include mki/base/operation_base.h #include mki_loader/op_register.h #include mki/utils/log/log.h #include atbops/params/params.h namespace Mki { using namespace Mki; static constexpr int32_t INPUT_NUM 2; static constexpr int32_t OUTPUT_NUM 1; class AddcustomOperation : public OperationBase { public: explicit AddcustomOperation(const std::string opName) noexcept : OperationBase(opName) {} /* ---------- 选择最佳kernel策略 ---------- */ Kernel *GetBestKernel(const LaunchParam launchParam) const override { // 具体函数实现请参考上面文件路径下的内容 } /* ---------- 张量数量 ---------- */ int64_t GetInputNum(const Any specificParam) const override { // 具体函数实现请参考上面文件路径下的内容 } int64_t GetOutputNum(const Any specificParam) const override { // 具体函数实现请参考上面文件路径下的内容 } protected: /* ---------- 形状推导 ---------- */ Status InferShapeImpl(const LaunchParam launchParam, SVectorTensor outTensors) const override { // 具体函数实现请参考上面文件路径下的内容 } }; /* ---------- 注册 ---------- */ REG_OPERATION(AddcustomOperation); } // namespace Mki构建与测试CMakeLists.txt文件路径example/ops/addcustom/CMakeLists.txt主要功能文件编译。set(addcustom_srcs ${CMAKE_CURRENT_LIST_DIR}/addcustom_operation.cpp ${CMAKE_CURRENT_LIST_DIR}/addcustom_kernel.cpp ${CMAKE_CURRENT_LIST_DIR}/tiling/addcustom_tiling.cpp ) add_operation(AddcustomOperation ${addcustom_srcs}) add_kernel(addcustom ascend910b vector op_kernel/addcustom.cpp AddcustomKernel)算子编译与环境变量设置Ascend Boost Comm仓的构建脚本文件为scripts/build.sh初次执行算子测试需要先编译测试框架testframeworkbash scripts/build.sh testframework随后编译example中的算子bash scripts/build.sh example编译后需要设置环境变量source output/mki/./set_env.sh测试为了对所编写的算子进行测试可在example/tests/pythontest/optest/目录下新增test_addcustom.pyimport os import unittest import numpy as np import torch import sys import logging sys.path.append(f{os.environ[MKI_HOME_PATH]}/tests/pythontest) import op_test OP_NAME AddcustomOperation OP_PARAM0 {addcustomDim: 0} class TestAddcustom(op_test.OpTest): def golden_calc(self, in_tensors): a in_tensors[0] b in_tensors[1] return [a b] def golden_compare(self, out_tensors, golden_out_tensors): return torch.allclose(out_tensors[0], golden_out_tensors[0], rtol0.001, atol0.001) def test_2d_half(self): shape (2 * 16,) a torch.randn(shape).to(torch.float16) b torch.randn(shape).to(torch.float16) self.set_param(OP_NAME, OP_PARAM0) self.execute([a, b], [torch.ones(shape).to(torch.float16)]) if __name__ __main__: unittest.main()通过以下命令执行测试脚本python example/tests/pythontest/optest/test_addcustom.py【免费下载链接】ascend-boost-comm算子公共平台南向对接不同组织开发的算子库北向支撑不同加速库应用实现M x N算子能力复用项目地址: https://gitcode.com/cann/ascend-boost-comm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考