基于Atlas进行项目-step3.1 算子开发-1：基于GPU，NPU等AI芯片实现自定义算子开发。

1.前情概括

在前两个博客中，我们已经详细介绍了如何使用Atlas 200 DK设备进行本地部署，并逐步掌握了深度学习模型的开发和基于本地NPU资源的计算加速。如果你对这些内容感兴趣，可以通过以下链接回顾并交流：

基于Atlas进行项目-step1 环境部署：Atlas 200 DK 环境配置，踩坑与学习记录_atlas200dk-CSDN博客

基于Atlas进行项目-step2 算法开发:在Atlas 200 DK（Soc=Ascend 310）快速上手自定义模型训练、部署与推理_aclliteresource-CSDN博客

然而，要充分利用NPU和GPU的计算能力，仅依赖如TensorFlow、PyTorch、MindSpore等深度学习框架来开发模型是不够的。这有两个主要原因：

• 广泛的计算需求：我们需要加速的不仅仅是AI模型的训练和推理，还有许多日常计算任务，例如密码加密与解密、数据压缩与通信，甚至包括万恶的“挖矿”程序。这些任务同样需要GPU、NPU等芯片的加速支持。然而，现有的算子库通常无法提供适用的接口，这就需要我们自行开发算子，申请GPU和NPU的资源来实现自定义算子，从而达到计算加速的目的。

• 算子库的局限性：现有框架以及CUDA、CANN自带的算子库功能有限，通常只支持神经网络中的特定计算方式的加速，例如矩阵乘法。由于GPU最初是为图像渲染而设计的，算子相对丰富一些；但NPU作为专为神经网络设计的AI芯片，特别是在国产自研的起步阶段，自带的算子相对较少。

因此，开发自定义算子不仅能方便我们解决实际问题，更重要的是，能够丰富国产自研AI生态，也算是为我国技术生态的建设贡献一份力量。这句话有点写项目书内味了哈。不过，本人也是希望国产的AI生态能够越来越好吧，可能会赶不上Nvidia的发展速度，但是能够有我们自己的AI生态，至少也不用整天担心芯片断供导致项目腰斩啥的了吧。

2.摘要

        近年来，随着人工智能（AI）和深度学习（DL）技术的快速发展，ChatGPT类自然语言大模型的普及与应用，比特币等区块链技术的崛起，以及如《黑悟空》这样的3A游戏大作的推出，高性能的GPU和NPU等硬件成为支撑这些应用的核心。可以说目前学术界，工业界，除了计算机学以外的生物学，医学，电力系统等众多学科也开始广泛依赖这些高性能计算资源。然而，对于许多非计算机专业人员，甚至是计算机领域但非嵌入式方向的从业者来说，利用GPU和NPU进行高效计算往往只能依赖如PyTorch、TensorFlow等框架进行机器学习任务。但是呢，学术要创新，工业讲性能，这两个指标仅依靠现有框架和算子库已经难以满足需求。想自己开发自定义算子呢，百度，Google翻半天也没个教程，问chatgpt呢（4o）,回答大概就是“这也太勾八难了，我只能给你一个大概指导，具体还是你自己搞吧”如下图所示。

为此，作者通过面向AI+博客+教材（《TBE 自定义算子开发指导_C75 TBE 自定义算子开发指导_C75》），探索出了一套基于GPU和NPU的自定义算子开发的完整流程。这套流程涵盖了从Linux环境配置（包括编译器版本的适配与安装）、Conda库的导入到代码编写的所有步骤，并通过一个简单的实例展示了成功完成自定义算子开发的过程。本文将首先介绍基于GPU的自定义算子开发，后续的博文将发布基于华为NPU的算子开发。

3.环境配置

首先作者本人的环境如下图所示：CUDA版本是12.4，驱动是550.54.15，nvcc版本是11.8。

看完上面的英伟达GPU配置，大家一定有个疑问就是 “你的显卡4090这么顶，为什么你CUDA不装12.6，nvcc还装11.X不装12的。”

说句实话，这些都是血的教训，是一次又一次的重装驱动，重装nvcc换来的后果，是自己摸索出来最适合目前这个环境，最稳定GPU开发配置。这些教训有以下几点。

1.如果你是LXD这种容器化程序来登录服务器的话，切记不要改变CUDA，驱动与nvcc的版本，一定不要，无论是升级也好降级也好。你一旦改了，不管成功与否，又或者后面又改回去了，你将永远失去使用GPU的权利，就连tensorflow这种自带算子库的都用不了了，报的错误就类似“检测到你的环境中有2套XXXX”就算和改之前的配置一样也会有问题，因为有一些什么依赖库残留，想删又删不掉，提示说是内核绑定，需要权限才能kill，遗憾的是，这个时候你是LXD的用户，没有权限，从此陷入了死局。

2.最好用conda来建立运行环境，不要再宿主设备上配置环境。最主要的原因就是上面那一条，一旦改变了CUDA和套件与驱动的版本之后，就会出现LXD容器外面映射的驱动信息不同，导致代码无法运行的问题。但是往往设备本身的CUDA与套件的版本要么都比较新（>12.X），要么就特别低，编译器也是，这个版本又必须改掉，这个时候最好的版本就是conda，并且实测，也是可行的。

3.就是想基于GPU和python开发自定义算子，需要用到这个pycuda这个库，而pytorch和tensorflow无法支持。比如，这个你想用GPU加速sha256这个计算函数，pytorch，tf自带库没有这个函数，这个时候如果想加速，就只能先把GPU的计算数据（xxxx.to(cpu)）取出来，放到CPU中然后用CPU的sha256()函数计算后，把结果又放回GPU的框架里面去。这个时候如果sha256正好是你论文创新点，是你做项目时候提升性能的关键，就非常的艹蛋了。于此同时这个pycuda这个库不知道是老了还是咋的，与nvcc>12.0的版本很不配合，经常在调用gcc的编译的时候，不是这个头文件丢失，就是那个XXX预定义名字不一样。因此这个需要将nvcc降到11.X（最好11.8）。顺便提一嘴，这个GCC的版本也必须降到11.X，10.X，9.X。如果这个时候你安装GCC是默认安装的话，（一般是13.X），在编译内核的时候也会报错。

总之，目前我摸索出来最稳定的配置就是，"CUDA版本是12.4，驱动是550.54.15，nvcc版本是11.8，gcc=11.4.0"

4.Conda配置

这里先用conda建立一个环境： Conda create --name env_name XXXX python=3.8（3.8，3.9，3.10都行）；然后conda activite env_name；然后安装pycudaXXXXX;具体的conda步骤，这里就不具体介绍了。

关键在于你激活之后需要注意，如下所示：

就是在你conda环境中安装nvcc。这个nvcc你一般安装pytorch或者tf的conda时，是不自带的，需要你自己安装的。这里有个坑，就是不要直接conda install -c conda-forge cudatoolkit
，这可能是AI告诉你的，也可能百度告诉你的，但是这种只能安装到最新版，且没办法通过=11.8.0这种改变版本。这里正确的命令应该是 conda install nvidia/label/cuda-11.8.0::cuda-nvcc。所有的版本命令可以在Cuda Nvcc | Anaconda.org的官网找到，具体如下所示：

To install this package run one of the following:
conda install nvidia::cuda-nvcc
conda install nvidia/label/cuda-11.3.0::cuda-nvcc
conda install nvidia/label/cuda-11.3.1::cuda-nvcc
conda install nvidia/label/cuda-11.4.0::cuda-nvcc
conda install nvidia/label/cuda-11.4.1::cuda-nvcc
conda install nvidia/label/cuda-11.4.2::cuda-nvcc
conda install nvidia/label/cuda-11.4.3::cuda-nvcc
conda install nvidia/label/cuda-11.4.4::cuda-nvcc
conda install nvidia/label/cuda-11.5.0::cuda-nvcc
conda install nvidia/label/cuda-11.5.1::cuda-nvcc
conda install nvidia/label/cuda-11.5.2::cuda-nvcc
conda install nvidia/label/cuda-11.6.0::cuda-nvcc
conda install nvidia/label/cuda-11.6.1::cuda-nvcc
conda install nvidia/label/cuda-11.6.2::cuda-nvcc
conda install nvidia/label/cuda-11.7.0::cuda-nvcc
conda install nvidia/label/cuda-11.7.1::cuda-nvcc
conda install nvidia/label/cuda-11.8.0::cuda-nvcc
conda install nvidia/label/cuda-12.0.0::cuda-nvcc
conda install nvidia/label/cuda-12.0.1::cuda-nvcc
conda install nvidia/label/cuda-12.1.0::cuda-nvcc
conda install nvidia/label/cuda-12.1.1::cuda-nvcc
conda install nvidia/label/cuda-12.2.0::cuda-nvcc
conda install nvidia/label/cuda-12.2.1::cuda-nvcc
conda install nvidia/label/cuda-12.2.2::cuda-nvcc
conda install nvidia/label/cuda-12.3.0::cuda-nvcc
conda install nvidia/label/cuda-12.3.1::cuda-nvcc
conda install nvidia/label/cuda-12.3.2::cuda-nvcc
conda install nvidia/label/cuda-12.4.0::cuda-nvcc
conda install nvidia/label/cuda-12.4.1::cuda-nvcc
conda install nvidia/label/cuda-12.5.0::cuda-nvcc
conda install nvidia/label/cuda-12.5.1::cuda-nvcc
conda install nvidia/label/cuda-12.6.0::cuda-nvcc

5.算子代码编写与解析

经过上面的环境配置后，就可以开始进入代码的编写了。下面是一段检测你是否具有调用本地GPU能力的检测程序：

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np

# 将 include_dirs 参数指向找到的 cuda_runtime.h 所在路径
mod = SourceModule("""
__global__ void test_kernel(float *a)
{
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  a[idx] = 2.0 * idx;
}
""", options=["-D__STRICT_ANSI__"],include_dirs=["/home/swb/anaconda3/envs/pow_env/lib/python3.9/site-packages/triton/backends/nvidia/include"])
# 准备测试数据
n = 256
a = np.zeros(n, dtype=np.float32)

# 分配GPU内存
a_gpu = cuda.mem_alloc(a.nbytes)

# 执行CUDA内核
test_kernel = mod.get_function("test_kernel")
test_kernel(a_gpu, block=(n, 1, 1))

# 将结果从GPU传回CPU
cuda.memcpy_dtoh(a, a_gpu)

# 输出结果
print(a)

这个代码分为4个部分。

1.头文件部分：

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np

这个没什么好说的，一般的都必须包括的，特别是这个SourceModule。

2.通过CUDA访问GPU的部分：

 将 include_dirs 参数指向找到的 cuda_runtime.h 所在路径
mod = SourceModule("""
__global__ void test_kernel(float *a)
{
  int idx = threadIdx.x + blockIdx.x * blockDim.x;
  a[idx] = 2.0 * idx;
}
""", options=["-D__STRICT_ANSI__"],include_dirs=["/home/swb/anaconda3/envs/pow_env/lib/python3.9/site-packages/triton/backends/nvidia/include"])
# 准备测试数据

这个地方又分为3个部分，第一个部分就是“”“ XXXXX""",这个地方就是你算子的主要组成部分，内核函数。通过C/C++语言编写，也就是用C/C++写一个函数 void XXX(){}，这个函数就可以理解是你用于调用GPU加速的,具体如何编写内容和逻辑呢，就靠你依据想实现的功能去实现了。上面这里的threadIdx.x，blockIdx.x，blockDim.x，在CUDA编程中，int idx = threadIdx.x + blockIdx.x * blockDim.x; 这一行代码的作用是计算每个线程在整个网格（grid）中的全局索引。

CUDA将线程组织成网格（grid）和线程块（block），具体有以下几个概念：

• threadIdx.x: 当前线程在其所在的线程块（block）内的索引。threadIdx.x 的值从0开始。
• blockIdx.x: 当前线程块在整个网格（grid）中的索引。blockIdx.x 的值从0开始。
• blockDim.x: 每个线程块中线程的总数，即线程块的维度。通常设定为线程块的大小。

这一行代码的含义是计算当前线程在整个网格（grid）中的全局索引 idx，即它在所有线程中的唯一标识。

• blockIdx.x * blockDim.x 计算了当前线程块（block）在整个网格中的起始位置。blockDim.x 是每个线程块中的线程数，而 blockIdx.x 是当前线程块的索引，因此 blockIdx.x * blockDim.x 是当前线程块的第一个线程在整个网格中的全局索引。

• threadIdx.x 是当前线程在其所在线程块中的索引，表示这个线程是它所在的线程块中的第几个线程。

因此，idx 代表了当前线程在整个网格（grid）中的全局索引位置。这个索引在许多CUDA程序中用于访问数组或其他数据结构，以确保每个线程处理不同的数据。（也就是并行的意义）下面提供另一段实例代码作为参考，助于理解：

# 完整代码
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np
import hashlib
import time

# 函数：使用 Python 计算 SHA-256 哈希值
def calculate_sha256(data, nonce):
    data_with_nonce = f"{data}{nonce}".encode()
    return hashlib.sha256(data_with_nonce).hexdigest()

# CUDA 核函数，用于计算哈希值并检查是否满足目标值
mod = SourceModule("""
#include <stdint.h>

__global__ void mine_kernel(char *data, uint64_t *nonce, unsigned long long *target, int *found, uint64_t *found_nonce) {
    uint64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    uint64_t local_nonce = *nonce + idx;

    // 将数据与 nonce 结合
    char local_data[80];
    memcpy(local_data, data, 76);  // 假设 data 长度为 76 字节
    memcpy(local_data + 76, &local_nonce, 4);

    // 计算哈希值
    uint32_t hash_value = 0x6a09e667;  // 这是一个简化的 SHA-256 模拟
    for (int i = 0; i < 80; i++) {
        hash_value ^= (hash_value << 5) + (hash_value >> 2) + (uint32_t)local_data[i];
    }

    // 比较哈希值和目标值
    if ((uint64_t)hash_value < *target && atomicExch(found, 1) == 0) {
        *found_nonce = local_nonce;
    }
}

""", options=["-D__STRICT_ANSI__"], include_dirs=["/home/swb/anaconda3/envs/pow_env/lib/python3.9/site-packages/triton/backends/nvidia/include"])

# 函数：在GPU上进行PoW挖矿
def gpu_mine(data, max_nonce, difficulty,kk):
    max_target = 2 ** 64 - 1
    target = min(2 ** (256 - difficulty), max_target)

    # 将目标值和其他变量传递给CUDA内核
    target_gpu = cuda.mem_alloc(8)
    cuda.memcpy_htod(target_gpu, np.array([target], dtype=np.uint64))

    found_gpu = cuda.mem_alloc(4)
    cuda.memcpy_htod(found_gpu, np.array([0], dtype=np.int32))

    found_nonce_gpu = cuda.mem_alloc(4)
    cuda.memcpy_htod(found_nonce_gpu, np.array([0], dtype=np.uint32))

    data_gpu = cuda.mem_alloc(len(data))
    cuda.memcpy_htod(data_gpu, np.array([ord(c) for c in data], dtype=np.uint8))

    mine_kernel = mod.get_function("mine_kernel")

    block_size = 256
    grid_size = max_nonce // block_size

    start_time = time.time()

    # 启动 CUDA 核函数
    mine_kernel(data_gpu, found_nonce_gpu, target_gpu, found_gpu, found_nonce_gpu,
                block=(block_size, 1, 1), grid=(grid_size, 1))

    cuda.Context.synchronize()

    # 从GPU获取结果
    found_nonce = np.zeros(1, dtype=np.uint32)
    cuda.memcpy_dtoh(found_nonce, found_nonce_gpu)

    end_time = time.time()

    # 验证找到的 nonce 是否满足难度要求
    if found_nonce[0] != 0:
        hash_value = calculate_sha256(data, found_nonce[0])
        if int(hash_value, 16) < target:
            print(f"成功挖掘到一个区块，nonce: {found_nonce[0]}")
            print(f"SHA-256 哈希值: {hash_value}")
        else:
            print(f"挖掘失败，nonce: {found_nonce[0]},哈希值: {hash_value}, 生成的哈希值不满足难度要求")
        print(f"耗时: {end_time - start_time} 秒")
    else:
        print("挖矿失败，未找到合适的 nonce 值。")

3.内核函数编译所需要依赖库地址：

 options=["-D__STRICT_ANSI__"], include_dirs=["/home/swb/anaconda3/envs/pow_env/lib/python3.9/site-packages/triton/backends/nvidia/include"]

这个地方，必须要有，要不然比报错，如下所示：

1. 缺少这个option或者你的gcc版本不对,就会报错 /usr/include/x86_64-linux-gnu/bits/mathcalls.h(62): error: identifier "_Float32" is undefined
2. 缺少后面那个include就会报错：“<command-line>: fatal error: cuda_runtime.h: No such file or directory compilation terminated.”

第一个的原因就是版本不匹配。第二个的原因就是因为conda和LXD的原因，这个库找依赖库会自动去\local\usr\cuda里面去找，但是conda和LXD的安装的nvcc和其他配置是是安装到了conda的文件夹里面，所以这个CUDA编译的依赖就会找不到，只能手动去找这个路径。

4.python代码调用内核函数使用GPU加速计算

n = 256
a = np.zeros(n, dtype=np.float32)

# 分配GPU内存
a_gpu = cuda.mem_alloc(a.nbytes)

# 执行CUDA内核
test_kernel = mod.get_function("test_kernel")
test_kernel(a_gpu, block=(n, 1, 1))

# 将结果从GPU传回CPU
cuda.memcpy_dtoh(a, a_gpu)

# 输出结果
print(a)

首先就是用cuda.mem_alloc()函数分配变量的内存，然后用get_function找到对应的你内核函数。接着把python代码从CPU定义变量传到GPU内核函数里面计算，完成GPU的加速。

大致流程就是这样。也就是说在自定义算子的过程中最重要的就是这个内核函数的设计，也就是用C/C++语言写多个函数，然后用python代码中将需要加速的方法和数据传到这个内核函数里面加速。

其实从代码的格式就可以看出来，mem_alloc和C的malloc基本是一样的，包括后面NPU的算子开发也一样，都是很多相似相通的地方，只要理解以后就一通百通了。 

最后，基于GPU进行算子开发的整体流程就如上所示啦。至于其他更进阶的方法和技巧，就后续在更新啦，或者在之后NPU博客中继续分享。今天时间有点晚，暂时的话就到这了，以上。