Mini-Infer (10): 卷积的终极形态 - Conv2D 实现与 BiasKernel 集成

1. 最后一块拼图：BiasKernel

在卷积操作 Output = Conv(Input, Weight) + Bias 中，加上偏置（Bias）是最后一步。虽然它计算量不大，但对于内存带宽要求很高。

为了保持架构的一致性，我们将 Bias 加法也封装为一个可调度、可优化的 Kernel。

bias.h: 接口定义

我们继续沿用 TensorRT 风格的注册表模式：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
// mini_infer/kernels/bias.h
namespace mini_infer {
namespace kernels {

// 定义函数签名：output += bias
template<typename T>
using BiasFunc = void(*)(T* output, const T* bias, int batch, int channel, int spatial);

// 定义注册表别名
DEFINE_REGISTRY_ALIAS(BiasRegistry, BiasFunc);

class BiasKernel {
public:
    // 统一入口：自动分发到最佳后端 (CPU/AVX/CUDA)
    template<typename T>
    static void add_channel_bias(
        T* output, const T* bias, 
        int batch_size, int channels, int spatial_size,
        KernelBackend backend = KernelBackend::AUTO
    ) {
        // 1. 初始化注册表
        KernelRegistryInitializer::initialize();
        
        // 2. 获取最佳内核 (AUTO 模式)
        auto func = (backend == KernelBackend::AUTO) 
            ? BiasRegistry<T>::instance().get_best_kernel()
            : BiasRegistry<T>::instance().get_kernel(backend);
        
        // 3. 执行
        if (func) func(output, bias, batch_size, channels, spatial_size);
        else throw std::runtime_error("No Bias kernel found!");
    }
    
    DEFINE_BACKEND_CHECKER(is_backend_available, BiasRegistry)
};

} // namespace kernels
} // namespace mini_infer

bias_cpu.cpp: CPU 实现

CPU 实现非常直观。注意这里的内存布局是 NCHW，所以偏置是加在 Channel 维度上的，这意味着对于同一个 Channel 的所有 spatial 像素（H*W），加的偏置值是相同的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// mini_infer/kernels/bias.cpp
template<typename T>
void bias_impl(T* output, const T* bias, int batch, int channels, int spatial) {
    for (int b = 0; b < batch; ++b) {
        T* batch_ptr = output + b * channels * spatial;
        for (int c = 0; c < channels; ++c) {
            T bias_val = bias[c];
            T* channel_ptr = batch_ptr + c * spatial;
            
            // 核心循环：这一步非常适合 SIMD 优化 (AVX2/AVX512)
            for (int s = 0; s < spatial; ++s) {
                channel_ptr[s] += bias_val;
            }
        }
    }
}

2. 核心战役：Conv2D 算子实现

有了 Im2Col、GEMM 和 Bias，Conv2D 的实现就不再是复杂的数学计算，而是一场优雅的数据流编排。

Conv2D 的设计选择：逐 Batch 处理

虽然“一次大 GEMM”听起来很酷，但它会导致内存布局错乱（CNHW），需要额外的转置操作。为了保持代码简洁且无需额外内存拷贝，我们选择逐 Batch 处理。

这意味着：

  1. 我们在 allocate_tensors 阶段就创建了正确的 NCHW 输出 Tensor。
  2. 在 forward 中，我们循环 Batch 维度 (N)。
  3. 在每次循环中，我们计算当前 Batch 的 output_n 指针偏移，并直接将 GEMM 结果写入其中。

这样，当循环结束时，输出内存天生就是完美的 NCHW 布局！

关键代码解析：forward

这是 Mini-Infer 目前最复杂的算子实现，请仔细阅读注释：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
// mini_infer/operators/conv2d.cpp

core::Status Conv2D::forward(...) {
    // ... (参数校验与 Tensor 获取) ...

    // 1. 准备数据指针
    const float* input_data = ...;
    const float* weight_data = ...;
    float* output_data = ...;

    // 2. 计算维度
    // M = Output Channels
    // K = Input Channels * KernelH * KernelW
    // N_gemm = Spatial Size (H_out * W_out)
    int M = C_out;
    int K = C_in * kernel_h * kernel_w;
    int N_gemm = H_out * W_out; // 单个 Batch 的 spatial size

    // 3. 准备 im2col 缓冲区 (Workspace)
    // 只需要为一个 Batch 申请内存，更好地利用 CPU 缓存 (L2/L3)
    int col_buffer_size = K * N_gemm;
    core::Buffer<float> col_buffer(col_buffer_size);

    // 4. 核心循环：逐 Batch 处理
    for (int n = 0; n < N; ++n) {
        // 计算当前 Batch 的指针偏移
        const float* input_n = input_data + n * C_in * H_in * W_in;
        float* output_n = output_data + n * C_out * N_gemm; // 直接指向 NCHW 中的正确位置
        
        // Step 1: im2col (将当前 Batch 的图片转化为矩阵)
        // Input: [C_in, H_in, W_in] -> Buffer: [K, N_gemm]
        kernels::Im2ColKernel::im2col<float>(
            input_n,
            col_buffer.data(),
            C_in, H_in, W_in,
            kernel_h, kernel_w,
            param_.stride_h, param_.stride_w,
            param_.padding_h, param_.padding_w,
            param_.dilation_h, param_.dilation_w,
            H_out, W_out
        );
        
        // Step 2: GEMM (计算卷积)
        // Weight (A): [M, K]
        // Buffer (B): [K, N_gemm]
        // Output (C): [M, N_gemm] <-- 这就是 [C_out, H_out*W_out]，即 NCHW 的一部分！
        kernels::GEMMKernel::gemm_nn<float>(
            weight_data,       // A
            col_buffer.data(), // B
            output_n,          // C (直接写入)
            M, N_gemm, K
        );
    }

    // 5. 执行 Bias 加法
    // Bias 加法支持批量处理，不需要写在循环里
    if (param_.use_bias) {
        kernels::BiasKernel::add_channel_bias<float>(
            output_data, bias_data,
            N, C_out, N_gemm
        );
    }

    return core::Status::SUCCESS;
}

💡 性能亮点：Cache Locality

虽然我们放弃了 Batched GEMM 的“一次性启动”，但我们获得了更好的缓存局部性 (Cache Locality)。

  • col_buffer 只需容纳一张图片的 im2col 结果。对于典型的 CNN 层，这个大小通常能放入 L2 或 L3 缓存中。
  • 我们频繁地复用这块“热”内存，而不是去访问一块巨大的冷内存。
  • 更重要的是，我们避免了额外的 Transpose 内存拷贝。

3. 总结与展望

至此，Mini-Infer 的核心功能开发阶段已经圆满结束！

回顾我们的成就：

  1. 架构：搭建了 Graph, Engine, Backend 分层架构。
  2. 内存：实现了 Allocator, Tensor, 以及 RAII Buffer。
  3. 内核：构建了基于模板元编程的 KernelRegistry，实现了 GEMM, Im2Col, Bias。
  4. 算子：实现了 ReLU, Linear, 以及最复杂的 Conv2D。