cuda-c-examples-torch

PyTorch + CUDA C 示例代码

本 Skill 包含完整的可运行示例代码，展示如何在 PyTorch 中使用 CUDA C 编写高性能 kernel，通过 load_inline JIT 编译集成。

集成模式

所有 CUDA C 内核都通过以下模式与 PyTorch 集成：

import torch
from torch.utils.cpp_extension import load_inline

# 1. CUDA 源代码（内核定义 + 调用函数）
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void my_kernel(const float* input, float* output, int size) {
    // 内核实现
}

torch::Tensor my_kernel_call(torch::Tensor input) {
    auto size = input.numel();
    auto output = torch::zeros_like(input);
    int block_size = 256;
    int num_blocks = (size + block_size - 1) / block_size;
    my_kernel<<<num_blocks, block_size>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), size);
    return output;
}
"""

# 2. C++ 函数声明
cpp_source = "torch::Tensor my_kernel_call(torch::Tensor input);"

# 3. JIT 编译
module = load_inline(
    name="my_cuda",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["my_kernel_call"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)

# 4. 调用
def my_op(x):
    return module.my_kernel_call(x)

示例列表

1. 向量加法（Vector Add）

算子类型: Element-wise 关键点:

• 最简单的 CUDA C 内核示例
• 一维索引和边界检查
• 标准五步模式

import torch
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void vector_add_kernel(
    const float* a, const float* b, float* c, int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

torch::Tensor vector_add_call(torch::Tensor a, torch::Tensor b) {
    auto n = a.numel();
    auto c = torch::empty_like(a);
    int block_size = 256;
    int num_blocks = (n + block_size - 1) / block_size;
    vector_add_kernel<<<num_blocks, block_size>>>(
        a.data_ptr<float>(), b.data_ptr<float>(),
        c.data_ptr<float>(), n);
    return c;
}
"""

cpp_source = "torch::Tensor vector_add_call(torch::Tensor a, torch::Tensor b);"

module = load_inline(
    name="vector_add_cuda",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vector_add_call"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)

def vector_add(a, b):
    return module.vector_add_call(a, b)