CUDA算子开发基础（一）

原创

已于 2025-07-10 10:25:27 修改 · 1.4k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

于 2025-06-17 11:43:23 首次发布

使用实现算子并加入torch

日常开发中我们经常能看到如何使用cuda编写算子，或者是cmake构建代码或者nvcc构建代码，cmake方式的优点在于可以管理更大规模的算子实现。这样可以复用一些基础代码，而nvcc的好处是方便调试同时最小依赖。两者都存在一个问题：我们真正实现一个算子往往是希望能在深度学习框架中调用它而不是单纯的实现算子。如何快速看到实现的算子在框架中使用也是算子开发中的重要一环。本系列文章尝试解决的核心问题是：

高效的算子开发基础组件，可复用的基础能力：

方便的单测
方便的性能测试
方便的可复用代码

方便的代码集成能力，可以轻松的在torch中调用并验证算子。
torch中集成高性能算子实现（triton）。

我们从第二个能力开始，它决定了你是否实际使用到这些算子而。有了这些嵌入框架算子能力之后我们甚至可以借助开源社区的最优代码实现扩充我们实际框架的算子能力。

实现一个简单的elu激活函数

首先我们需要知道elu函数是如何计算的：
$\begin{cases} x & \quad x \geq 0 \\ \alpha\cdot(e^x - 1) &\quad x < 0 \end{cases}$
其中， $α\alpha$ 为超参数，通常设为1。

这是一个简单地elementwise算子。

通过CUDA_EXTENTION+PyBind构造自定义算子

.
├── include
│   └── elu.h
├── main.py
├── model.py
├── setup.py
└── src
    ├── elu_bind.cc
    ├── elu.cc
    └── elu.cu

先声明头文件include/elu.h

#pragma once
#include <torch/extension.h>
torch::Tensor elu_cuda(const torch::Tensor &a);
torch::Tensor elu_cpu(const torch::Tensor &a);

上面的代码声明了对CPU和GPU的接口，然后我们将在C++代码和cu代码中分别实现对应的逻辑。首先是src/elu.cu的cuda实现：

#include "elu.h"
#include <cooperative_groups.h>
#include <cuda_runtime.h>
#include <torch/extension.h>
#define ALPHA 1.f
__device__ __forceinline__ float elu_fp32(float x) {
   
   
    return x > 0 ? x : ALPHA * (__expf(x) - 1);
}
__global__ void elu_kernel_fp32(float *d_out, const float *d_in, const int N) {
   
   
    extern __shared__ float smem[