PyTorch框架

通过PyTorch框架进行模型的训练、推理时，会调用很多算子进行计算。开发者开发的自定义算子如果需要集成部署到PyTorch框架，有如下几种方式：

Kernel直调：通过适配Pybind调用，可以实现PyTorch框架调用算子Kernel程序。
单算子API调用：该模式下的适配插件开发流程和具体样例请参见《Ascend Extension for PyTorch 框架特性指南》中的“基于OpPlugin算子适配开发”章节。
图模式调用：自定义算子在Pytorch图模式下的适配开发指导请参见《PyTorch图模式使用指南(TorchAir)》中的“自定义算子入图”章节。

图1 Pytorch框架部署方式

本节主要提供通过Pybind实现PyTorch框架调用算子Kernel程序的指导。

Pybind调用介绍

Pybind是一个用于将C++代码与Python解释器集成的库，实现原理是通过将C++代码编译成动态链接库（DLL）或共享对象（SO）文件，使用Pybind提供的API将算子核函数与Python解释器进行绑定。在Python解释器中使用绑定的C++函数、类和变量，从而实现Python与C++代码的交互。

环境准备

基于环境准备，还需要安装以下依赖：

安装PyTorch框架
安装torch_npu插件
安装pybind11
```
 pip3 install pybind11
```

工程目录

样例工程目录结构如下所示：

├── add_custom_test.py      // Python调用脚本 
├── add_custom.asc          // 算子实现 + pybind11函数封装 
└── CMakeLists.txt          // 编译工程文件

基于该算子工程，开发者进行算子开发的步骤如下：

完成算子Kernel侧实现；编写算子调用应用程序并定义pybind模块。这两部分代码均在add_custom.asc中实现。
编写Python调用脚本add_custom_test.py，包括生成输入数据和真值数据，调用封装的模块以及验证结果。
编写CMake编译配置文件CMakeLists.txt。
完成算子的编译运行和结果验证。

算子Kernel实现、调用和Pybind模块定义

下面代码以add_custom算子为例，介绍算子核函数实现及调用的应用程序add_custom.asc如何编写。您在实现自己的应用程序时，需要关注由于算子核函数不同带来的修改，包括算子核函数名，入参出参的不同等，相关API的调用方式直接复用即可。同时该文件中定义了Pybind模块，便于后续在Python脚本中调用。

按需包含头文件。

        
             // Pybind和PyTorch调用所需的头文件
#include <pybind11/pybind11.h>
#include <torch/extension.h>

#include "torch_npu/csrc/core/npu/NPUStream.h"
// Kernel侧实现需要的头文件
#include "kernel_operator.h"

算子Kernel实现。

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
    {
        this->blockLength = totalLength / AscendC::GetBlockNum();
        this->tileNum = 8;
        this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;
        xGm.SetGlobalBuffer((__gm__ half *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        yGm.SetGlobalBuffer((__gm__ half *)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        zGm.SetGlobalBuffer((__gm__ half *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(half));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        AscendC::LocalTensor<half> yLocal = inQueueY.AllocTensor<half>();
        AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        AscendC::LocalTensor<half> yLocal = inQueueY.DeQue<half>();
        AscendC::LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
        AscendC::Add(zLocal, xLocal, yLocal, this->tileLength);
        outQueueZ.EnQue<half>(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<half> zLocal = outQueueZ.DeQue<half>();
        AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<half> xGm;
    AscendC::GlobalTensor<half> yGm;
    AscendC::GlobalTensor<half> zGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
};

__global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    KernelAdd op;
    op.Init(x, y, z, totalLength);
    op.Process();
}

算子调用程序编写，使用<<<>>>接口调用算子核函数完成指定的运算。样例中的c10_npu::getCurrentNPUStream接口用于获取当前npu流，返回值类型NPUStream，使用方式请参考《Ascend Extension for PyTorch 自定义API参考》中的“（beta）c10_npu::getCurrentNPUStream”章节。

需要注意的是，本样例的输入x，y的内存是在Python调用脚本add_custom_test.py中分配的。

         
          
            
            
              namespace my_add {
at::Tensor run_add_custom(const at::Tensor &x, const at::Tensor &y)
{
   // 运行资源申请，通过c10_npu::getCurrentNPUStream()的函数获取当前NPU上的流
    auto aclStream = c10_npu::getCurrentNPUStream().stream(false);
     // 分配Device侧输出内存
    at::Tensor z = at::empty_like(x);
    uint32_t blockDim = 8;
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }
    // 用<<<>>>接口调用核函数完成指定的运算
    auto xGm = static_cast<uint8_t *>(const_cast<void *>(x.storage().data()));
    auto yGm = static_cast<uint8_t *>(const_cast<void *>(y.storage().data()));
    auto zGm = static_cast<uint8_t *>(const_cast<void *>(z.storage().data()));
    add_custom<<<blockDim, nullptr, aclStream>>>(xGm, yGm, zGm, totalLength);
    // 将Device上的运算结果拷贝回Host并释放申请的资源
    return z;
}
} // namespace my_add

             

           

         
        

定义Pybind模块，将C++函数封装成Python函数。PYBIND11_MODULE是Pybind11库中的一个宏，用于定义一个Python模块。它接受两个参数，第一个参数是封装后的模块名，第二个参数是一个Pybind11模块对象，用于定义模块中的函数、类、常量等。通过调用m.def()方法，可以将步骤3中函数my_add::run_add_custom()转成Python函数run_add_custom，使其可以在Python代码中被调用。

        
             PYBIND11_MODULE(add_custom, m) { // 模块名add_custom，模块对象m
  m.doc() = "add_custom pybind11 interfaces";  // optional module docstring
  m.def("run_add_custom", &my_add::run_add_custom, ""); // 将函数run_add_custom与Pybind模块进行绑定
}

Python调用脚本

在Python调用脚本中，使用torch接口生成随机输入数据并分配内存，通过导入封装的自定义模块add_custom，调用自定义模块add_custom中的run_add_custom函数，从而在NPU上执行算子。算子核函数NPU侧运行验证的步骤如图2。

图2 NPU侧运行验证原理图

import sys
import os
import torch
import torch_npu
from torch_npu.testing.testcase import TestCase, run_tests
sys.path.append(os.getcwd())
import add_custom
torch.npu.config.allow_internal_format = False
class TestCustomAdd(TestCase):
    def test_add_custom_ops(self):
        # 分配Host侧输入内存，并进行数据初始化
        length = [8, 2048]
        x = torch.rand(length, device='cpu', dtype=torch.float16)
        y = torch.rand(length, device='cpu', dtype=torch.float16)
        # 分配Device侧输入内存，并将数据从Host上拷贝到Device上
        x_npu = x.npu()
        y_npu = y.npu()
        output = add_custom.run_add_custom(x_npu, y_npu)
        cpuout = torch.add(x, y)
        self.assertRtolEqual(output, cpuout)
if __name__ == "__main__":
    run_tests()

CMake编译配置文件编写

该示例中通过CMake脚本生成算子对应的动态库后，Python中会通过import加载该动态库后执行计算。如果需要了解算子编译更多内容，请参考通过CMake编译。

cmake_minimum_required(VERSION 3.16)

find_package(ASC REQUIRED HINTS $ENV{ASCEND_INSTALL_PATH}/compiler/tikcpp/ascendc_kernel_cmake)

execute_process(COMMAND python3 -c "import os; import torch; print(os.path.dirname(torch.__file__))"
    OUTPUT_STRIP_TRAILING_WHITESPACE
    OUTPUT_VARIABLE TORCH_PATH
)
message("TORCH_PATH is ${TORCH_PATH}")

execute_process(COMMAND python3 -c "import os; import torch_npu; print(os.path.dirname(torch_npu.__file__))"
    OUTPUT_STRIP_TRAILING_WHITESPACE
    OUTPUT_VARIABLE TORCH_NPU_PATH
)
message("TORCH_NPU_PATH is ${TORCH_NPU_PATH}")

execute_process(COMMAND python3 -m pybind11 --includes
    OUTPUT_STRIP_TRAILING_WHITESPACE
    OUTPUT_VARIABLE PYBIND11_INC
)
string(REPLACE " " ";" PYBIND11_INC ${PYBIND11_INC})

execute_process(COMMAND python3-config --extension-suffix
    OUTPUT_STRIP_TRAILING_WHITESPACE
    OUTPUT_VARIABLE PYBIND11_SUFFIX
)

project(kernel_samples LANGUAGES ASC CXX)

add_library(pybind11_lib SHARED
    add_custom.asc
)

target_link_libraries(pybind11_lib PRIVATE
    torch_npu
)

target_link_directories(pybind11_lib PRIVATE
    ${TORCH_PATH}/lib
    ${TORCH_NPU_PATH}/lib
)

target_include_directories(pybind11_lib PRIVATE
    ${TORCH_NPU_PATH}/include
    ${TORCH_PATH}/include
    ${TORCH_PATH}/include/torch/csrc/api/include
)

target_compile_definitions(pybind11_lib PRIVATE
    _GLIBCXX_USE_CXX11_ABI=0
)

target_compile_options(pybind11_lib PRIVATE
    ${PYBIND11_INC}
    $<$<COMPILE_LANGUAGE:ASC>:--npu-arch=dav-2201>
    -fPIC
)

set_target_properties(pybind11_lib PROPERTIES
    OUTPUT_NAME add_custom${PYBIND11_SUFFIX}
    PREFIX "" SUFFIX ""
)

编译和运行程序

完成上述文件的编写后，可以执行如下命令编译和运行应用程序。

rm -rf build; mkdir -p build; cd build  # 创建并进入build目录
cmake ..; make -j                       # 编译算子so
python3 ../add_custom_test.py           # 执行样例

父主题： AI框架算子适配