1
#include <torch/extension.h>
2
#include <cuda_runtime.h>
3

4
// 高性能向量加法 kernel (float4 向量化)
5
__global__ void add_kernel_vec4(const float* __restrict__ a,
6
                                const float* __restrict__ b,
7
                                float* __restrict__ c,
8
                                int N) {
9
    const float4* a4 = reinterpret_cast<const float4*>(a);
10
    const float4* b4 = reinterpret_cast<const float4*>(b);
11
    float4* c4 = reinterpret_cast<float4*>(c);
12

13
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
14
    int stride = blockDim.x * gridDim.x;
15
    int N4 = N / 4;
16

17
    for (int i = idx; i < N4; i += stride) {
18
        float4 av = a4[i];
19
        float4 bv = b4[i];
20
        float4 cv;
21
        cv.x = av.x + bv.x;
22
        cv.y = av.y + bv.y;
23
        cv.z = av.z + bv.z;
24
        cv.w = av.w + bv.w;
25
        c4[i] = cv;
26
    }
27

28
    int remainder_start = N4 * 4;
29
    for (int i = remainder_start + idx; i < N; i += stride) {
30
        c[i] = a[i] + b[i];
31
    }
32
}
33

34
torch::Tensor add_forward(torch::Tensor a, torch::Tensor b) {
35
    TORCH_CHECK(a.device().is_cuda(), "a must be CUDA tensor");
36
    TORCH_CHECK(b.device().is_cuda(), "b must be CUDA tensor");
37
    auto c = torch::empty_like(a);
38
    int N = a.numel();
39
    const int threads = 256;
40
    int blocks = ((N / 4) + threads - 1) / threads;
41
    if (blocks == 0) blocks = 1;
42
    add_kernel_vec4<<<blocks, threads>>>(
43
        a.data_ptr<float>(),
44
        b.data_ptr<float>(),
45
        c.data_ptr<float>(),
46
        N
47
    );
48
    return c;
49
}
50

51
// 自动微分包装
52
class AddFunction : public torch::autograd::Function<AddFunction> {
53
public:
54
    static torch::Tensor forward(
55
        torch::autograd::AutogradContext* ctx,
56
        torch::Tensor a,
57
        torch::Tensor b) {
58
        ctx->save_for_backward({a, b});
59
        return add_forward(a, b);
60
    }
61

62
    static torch::autograd::variable_list backward(
63
        torch::autograd::AutogradContext* ctx,
64
        torch::autograd::variable_list grad_output) {
65
        auto grad = grad_output[0];
66
        return {grad, grad};
67
    }
68
};
69

70
torch::Tensor add_autograd(torch::Tensor a, torch::Tensor b) {
71
    return AddFunction::apply(a, b);
72
}
73

74
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
75
    m.def("add", &add_autograd, "Vector addition with autograd");
76
    m.def("add_forward", &add_forward, "Raw forward (no autograd)");
77
}

1
from setuptools import setup
2
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
3

4
setup(
5
    name="cuda_add",
6
    version="0.1.0",
7
    ext_modules=[
8
        CUDAExtension(
9
            name="cuda_add",
10
            sources=["vector_add.cu"],
11
            extra_compile_args={"cxx": ["-O3"], "nvcc": ["-O3"]}
12
        )
13
    ],
14
    cmdclass={"build_ext": BuildExtension},
15
    install_requires=["torch"],
16
)

1
#include <torch/extension.h>
2
#include <cuda_runtime.h>
3

4
// 定义 uchar3 结构体（与 OpenCV 对应）
5
struct uchar3 {
6
    unsigned char x, y, z;
7
};
8

9
__global__ void rgbToGrayKernel(const uchar3* img, unsigned char* gray,
10
                                int width, int height) {
11
    int x = blockIdx.x * blockDim.x + threadIdx.x;
12
    int y = blockIdx.y * blockDim.y + threadIdx.y;
13
    if (x < width && y < height) {
14
        int idx = y * width + x;
15
        uchar3 pixel = img[idx];
16
        gray[idx] = (unsigned char)(0.299f * pixel.x + 0.587f * pixel.y + 0.114f * pixel.z);
17
    }
18
}
19

20
torch::Tensor rgb_to_gray_cuda(torch::Tensor img) {
21
    // 假设输入是 uint8 [H, W, 3] 且已在 GPU 上
22
    TORCH_CHECK(img.dim() == 3 && img.size(2) == 3, "Input must be HxWx3");
23
    TORCH_CHECK(img.dtype() == torch::kUInt8, "Input must be uint8");
24
    int height = img.size(0);
25
    int width = img.size(1);
26
    auto gray = torch::empty({height, width}, img.options().dtype(torch::kUInt8));
27

28
    dim3 block(16, 16);
29
    dim3 grid((width + 15) / 16, (height + 15) / 16);
30
    rgbToGrayKernel<<<grid, block>>>(
31
        reinterpret_cast<const uchar3*>(img.data_ptr<unsigned char>()),
32
        gray.data_ptr<unsigned char>(),
33
        width, height
34
    );
35
    return gray;
36
}
37

38
__global__ void normalizeVolumeKernel(const unsigned short* in, float* out,
39
                                      int dimX, int dimY, int dimZ, float maxVal) {
40
    int x = blockIdx.x * blockDim.x + threadIdx.x;
41
    int y = blockIdx.y * blockDim.y + threadIdx.y;
42
    int z = blockIdx.z * blockDim.z + threadIdx.z;
43
    if (x < dimX && y < dimY && z < dimZ) {
44
        int idx = z * dimY * dimX + y * dimX + x;
45
        out[idx] = (float)in[idx] / maxVal;
46
    }
47
}
48

49
torch::Tensor normalize_volume_cuda(torch::Tensor volume, float maxVal) {
50
    TORCH_CHECK(volume.dim() == 3, "Volume must be 3D");
51
    TORCH_CHECK(volume.dtype() == torch::kUInt16, "Input must be uint16");
52
    int dimX = volume.size(2);  // 注意 PyTorch 维度顺序 DxHxW，这里假设 Z,Y,X
53
    int dimY = volume.size(1);
54
    int dimZ = volume.size(0);
55
    auto out = torch::empty_like(volume, volume.options().dtype(torch::kFloat32));
56

57
    dim3 block(8, 8, 4);
58
    dim3 grid((dimX + 7) / 8, (dimY + 7) / 8, (dimZ + 3) / 4);
59
    normalizeVolumeKernel<<<grid, block>>>(
60
        volume.data_ptr<unsigned short>(),
61
        out.data_ptr<float>(),
62
        dimX, dimY, dimZ, maxVal
63
    );
64
    return out;
65
}
66

67
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
68
    m.def("rgb_to_gray", &rgb_to_gray_cuda, "RGB to grayscale (CUDA)");
69
    m.def("normalize_volume", &normalize_volume_cuda, "3D volume normalization (CUDA)");
70
}

1
import torch
2
import cv2
3
import numpy as np
4
import custom_image_ops  # 编译后的模块名
5

6
# RGB 转灰度测试
7
img = cv2.imread("input.jpg")
8
img_tensor = torch.from_numpy(img).cuda()  # HxWx3
9
gray = custom_image_ops.rgb_to_gray(img_tensor)
10
cv2.imwrite("gray.png", gray.cpu().numpy())
11

12
# 三维归一化测试
13
vol = torch.randint(0, 4096, (128, 256, 256), dtype=torch.uint16, device='cuda')
14
norm = custom_image_ops.normalize_volume(vol, 4095.0)
15
print(norm.min(), norm.max())  # 0.0 ~ 1.0