cpu-basics

CPU C++ 编程基础

1. 核心概念

内核 (Kernel)

• 定义: 使用 PYBIND11_MODULE 注册的 C++ 函数，编译后在 CPU 上执行
• 特点: 直接操作张量数据指针，支持多种数据类型
• 形式: 使用 PyTorch C++ 扩展，通过 load_inline 动态编译加载

张量处理

• 连续性: 确保张量内存布局连续，避免非连续访问
• 类型统一: 内部计算使用统一类型（优先 float32/float64/int32/int64），最后转换回原类型
• 边界检查: 所有数组访问前必须检查边界

内存管理

• 自动管理: PyTorch 自动管理张量内存生命周期
• 指针操作: 直接操作数据指针进行高效计算
• 类型安全: 确保指针类型与张量类型匹配

2. 标准内核结构（五步模式）

所有 CPU C++ 内核都遵循相同的五步结构模式：

torch::Tensor standard_kernel(torch::Tensor x) {
    // 1. 确保输入张量是连续的
    if (!x.is_contiguous()) {
        x = x.contiguous();
    }
    
    // 2. 检查数据类型，支持多种类型
    torch::ScalarType dtype = x.scalar_type();
    bool need_convert = (dtype != torch::kFloat32 && dtype != torch::kFloat64 && 
                        dtype != torch::kInt32 && dtype != torch::kInt64);
    torch::Tensor input = need_convert ? x.to(torch::kFloat32) : x;

    // 3. 创建输出张量
    torch::Tensor output = torch::zeros_like(input);

    // 4. 根据数据类型分发计算
    if (input.scalar_type() == torch::kFloat32) {
        auto x_ptr = input.data_ptr<float>();
        auto out_ptr = output.data_ptr<float>();
        int64_t numel = input.numel();
        for (int64_t i = 0; i < numel; ++i) {
            out_ptr[i] = std::max(0.0f, x_ptr[i]);  // ReLU: max(0, x)
        }
    } else if (input.scalar_type() == torch::kFloat64) {
        auto x_ptr = input.data_ptr<double>();
        auto out_ptr = output.data_ptr<double>();
        int64_t numel = input.numel();
        for (int64_t i = 0; i < numel; ++i) {
            out_ptr[i] = std::max(0.0, x_ptr[i]);
        }
    } else if (input.scalar_type() == torch::kInt32) {
        auto x_ptr = input.data_ptr<int32_t>();
        auto out_ptr = output.data_ptr<int32_t>();
        int64_t numel = input.numel();
        for (int64_t i = 0; i < numel; ++i) {
            out_ptr[i] = std::max(0, x_ptr[i]);
        }
    } else if (input.scalar_type() == torch::kInt64) {
        auto x_ptr = input.data_ptr<int64_t>();
        auto out_ptr = output.data_ptr<int64_t>();
        int64_t numel = input.numel();
        for (int64_t i = 0; i < numel; ++i) {
            out_ptr[i] = std::max(0L, x_ptr[i]);
        }
    }

    // 5. 转换回原类型
    if (need_convert) {
        output = output.to(dtype);
    }
    return output;
}