如果已經通過Ascend C編程語言實現了算子，那該如何通過pybind進行調用呢踩官？

1 Pybind調用介紹

通過PyTorch框架進行模型的訓練惧互、推理時，會調用很多算子進行計算占遥，其中的調用方式與kernel編譯流程有關。

• 對于自定義算子工程输瓜，需要使用PyTorch Ascend Adapter中的OP-Plugin算子插件對功能進行擴展瓦胎，讓torch可以直接調用自定義算子包中的算子，詳細內容可以參考PyTorch框架尤揣；
• 對于KernelLaunch開放式算子編程的方式搔啊，通過適配

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

2 工程目錄結構

該樣例的工程目錄結構如下：

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

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

• 完成算子kernel側實現。
• 編寫算子調用應用程序和定義pybind模塊pybind11.cpp陪拘。
• 編寫Python調用腳本add_custom_test.py厂镇，包括生成輸入- 數據和真值數據，調用封裝的模塊以及驗證結果左刽。
• 編寫CMake編譯配置文件CMakeLists.txt捺信。
• 根據實際需要修改編譯運行算子的腳本run.sh并執(zhí)行該腳本，完成算子的編譯運行和結果驗證欠痴。

3 環(huá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 工程實現

4.1 算子kernel實現

之前的文章中掌挚，已經實現過，add_custom.cpp內容如下：

/**
 * @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 實現pybind11.cpp

1菩咨、按需包含頭文件吠式。
需要注意的是，需要包含對應的核函數調用接口聲明所在的頭文件alcrtlaunch_{kernel_name}.h（該頭文件為工程框架自動生成抽米，

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

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

2云茸、編寫框架調用程序

at::Tensor run_add_custom(const at::Tensor &x, const at::Tensor &y)
{
    // 運行資源申請是目，通過c10_npu::getCurrentNPUStream()的函數獲取當前NPU上的流
    auto acl_stream = 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;
    }
    // 用ACLRT_LAUNCH_KERNEL接口調用核函數完成指定的運算
    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上的運算結果拷貝回Host并釋放申請的資源
     return z;
}

需要注意的是，輸入x查辩，y的內存是在Python調用腳本add_custom_test.py（往下看）中分配的胖笛。
3、 定義Pybind模塊
將C++函數封裝成Python函數宜岛。PYBIND11_MODULE是Pybind11庫中的一個宏，用于定義一個Python模塊功舀。它接受兩個參數萍倡，第一個參數是封裝后的模塊名，第二個參數是一個Pybind11模塊對象辟汰，用于定義模塊中的函數列敲、類阱佛、常量等。通過調用m.def()方法戴而，可以將步驟2中函數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模塊進行綁定
}

4.3 編寫Python調用腳本

在Python調用腳本中淮逊，使用torch接口生成隨機輸入數據并分配內存，通過導入封裝的自定義模塊add_custom扶踊，調用自定義模塊add_custom中的run_add_custom函數泄鹏，從而在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側輸入內存秧耗，并進行數據初始化
        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()

4.4 編寫CMakeLists實現pybind11文件編譯

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