AscendC從入門到精通系列（四）使用Pybind調(diào)用AscendC算子

如果已經(jīng)通過Ascend C編程語言實現(xiàn)了算子，那該如何通過pybind進行調(diào)用呢？

1 Pybind調(diào)用介紹

通過PyTorch框架進行模型的訓(xùn)練、推理時，會調(diào)用很多算子進行計算，其中的調(diào)用方式與kernel編譯流程有關(guān)。

  • 對于自定義算子工程，需要使用PyTorch Ascend Adapter中的OP-Plugin算子插件對功能進行擴展，讓torch可以直接調(diào)用自定義算子包中的算子，詳細(xì)內(nèi)容可以參考PyTorch框架；
  • 對于KernelLaunch開放式算子編程的方式，通過適配

Pybind調(diào)用，可以實現(xiàn)PyTorch框架調(diào)用算子kernel程序。
Pybind是一個用于將C++代碼與Python解釋器集成的庫，實現(xiàn)原理是通過將C++代碼編譯成動態(tài)鏈接庫（DLL）或共享對象（SO）文件，使用Pybind提供的API將算子核函數(shù)與Python解釋器進行綁定。在Python解釋器中使用綁定的C++函數(shù)、類和變量，從而實現(xiàn)Python與C++代碼的交互。在Kernel直調(diào)中使用時，就是將Pybind模塊與算子核函數(shù)進行綁定，將其封裝成Python模塊，從而實現(xiàn)兩者交互。

2 工程目錄結(jié)構(gòu)

該樣例的工程目錄結(jié)構(gòu)如下：

├── CppExtensions 
│   ├── add_custom_test.py      // Python調(diào)用腳本 
│   ├── add_custom.cpp          // 算子實現(xiàn) 
│   ├── CMakeLists.txt          // 編譯工程文件 
│   ├── pybind11.cpp            // pybind11函數(shù)封裝
│   └── run.sh                  // 編譯運行算子的腳本

基于該算子工程，開發(fā)者進行算子開發(fā)的步驟如下：

  • 完成算子kernel側(cè)實現(xiàn)。
  • 編寫算子調(diào)用應(yīng)用程序和定義pybind模塊pybind11.cpp。
  • 編寫Python調(diào)用腳本add_custom_test.py，包括生成輸入- 數(shù)據(jù)和真值數(shù)據(jù)，調(diào)用封裝的模塊以及驗證結(jié)果。
  • 編寫CMake編譯配置文件CMakeLists.txt。
  • 根據(jù)實際需要修改編譯運行算子的腳本run.sh并執(zhí)行該腳本，完成算子的編譯運行和結(jié)果驗證。

3 環(huán)境準(zhǔn)備

3.1安裝pytorch (這里以2.1.0版本為例)

// aarch64環(huán)境上安裝
pip3 install torch==2.1.0

// x86環(huán)境上安裝
pip3 install torch==2.1.0+cpu  --index-url https://download.pytorch.org/whl/cpu

3.2 安裝torch-npu（昇騰適配torch的開發(fā)工程，這里以Pytorch2.1.0、python3.9、CANN版本8.0.RC1.alpha002為例）

 git clone https://gitee.com/ascend/pytorch.git -b v6.0.rc1.alpha002-pytorch2.1.0
 cd pytorch/
 bash ci/build.sh --python=3.9
 pip3 install dist/*.whl

3.3 安裝pybind11

pip3 install pybind11

4 工程實現(xiàn)

4.1 算子kernel實現(xiàn)

之前的文章中，已經(jīng)實現(xiàn)過，add_custom.cpp內(nèi)容如下：

/**
 * @file add_custom.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "kernel_operator.h"
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::QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::QuePosition::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;
};

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

4.2 實現(xiàn)pybind11.cpp

1、按需包含頭文件。
需要注意的是，需要包含對應(yīng)的核函數(shù)調(diào)用接口聲明所在的頭文件alcrtlaunch_{kernel_name}.h（該頭文件為工程框架自動生成，

#include"aclrtlaunch_add_custom.h"），kernel_name為算子核函數(shù)的名稱。
#include <pybind11/pybind11.h>
#include <torch/extension.h>

#include "aclrtlaunch_add_custom.h"
#include "torch_npu/csrc/core/npu/NPUStream.h"

2、編寫框架調(diào)用程序

at::Tensor run_add_custom(const at::Tensor &x, const at::Tensor &y)
{
    // 運行資源申請，通過c10_npu::getCurrentNPUStream()的函數(shù)獲取當(dāng)前NPU上的流
    auto acl_stream = c10_npu::getCurrentNPUStream().stream(false);
    // 分配Device側(cè)輸出內(nèi)存
    at::Tensor z = at::empty_like(x);
    uint32_t blockDim = 8;
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }
    // 用ACLRT_LAUNCH_KERNEL接口調(diào)用核函數(shù)完成指定的運算
    ACLRT_LAUNCH_KERNEL(add_custom)
    (blockDim, acl_stream, const_cast<void *>(x.storage().data()), const_cast<void *>(y.storage().data()),
     const_cast<void *>(z.storage().data()), totalLength);
     // 將Device上的運算結(jié)果拷貝回Host并釋放申請的資源
     return z;
}

需要注意的是，輸入x，y的內(nèi)存是在Python調(diào)用腳本add_custom_test.py（往下看）中分配的。
3、 定義Pybind模塊
將C++函數(shù)封裝成Python函數(shù)。PYBIND11_MODULE是Pybind11庫中的一個宏，用于定義一個Python模塊。它接受兩個參數(shù)，第一個參數(shù)是封裝后的模塊名，第二個參數(shù)是一個Pybind11模塊對象，用于定義模塊中的函數(shù)、類、常量等。通過調(diào)用m.def()方法，可以將步驟2中函數(shù)my_add::run_add_custom()轉(zhuǎn)成Python函數(shù)run_add_custom，使其可以在Python代碼中被調(diào)用。

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, ""); // 將函數(shù)run_add_custom與Pybind模塊進行綁定
}

4.3 編寫Python調(diào)用腳本

在Python調(diào)用腳本中，使用torch接口生成隨機輸入數(shù)據(jù)并分配內(nèi)存，通過導(dǎo)入封裝的自定義模塊add_custom，調(diào)用自定義模塊add_custom中的run_add_custom函數(shù)，從而在NPU上執(zhí)行算子。

import torch
import torch_npu
from torch_npu.testing.testcase import TestCase, run_tests
import sys, os
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側(cè)輸入內(nèi)存，并進行數(shù)據(jù)初始化
        length = [8, 2048]
        x = torch.rand(length, device='cpu', dtype=torch.float16)
        y = torch.rand(length, device='cpu', dtype=torch.float16)
        // 分配Device側(cè)輸入內(nèi)存，并將數(shù)據(jù)從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()

4.4 編寫CMakeLists實現(xiàn)pybind11文件編譯

編譯進工程的方式有很多，各個項目不一樣，這里提供一個參考：
operator/AddCustomSample/KernelLaunch/CppExtensions/CMakeLists.txt · Ascend/samples - 碼云 - 開源中國 (gitee.com)

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容