3D Convolution

原始题目：LeetGPU - 3D Convolution

题目描述

编写一个 GPU 程序，执行三维卷积操作。给定三维输入体积和三维卷积核，使用 "valid" 边界条件（无填充）计算卷积输出。

对于三维卷积，位置 $(i, j, k)$ 的输出为：

$$output(i, j, k) = \sum_{d=0}^{K_d-1}\sum_{r=0}^{K_r-1}\sum_{c=0}^{K_c-1} input(i+d, j+r, k+c) \cdot kernel(d, r, c)$$

所有数据（输入、卷积核、输出）均以一维数组表示，存储顺序为：先行列（row-major），再深度。

输出维度为：

$$\begin{aligned} output\_depth &= input\_depth - kernel\_depth + 1 \\ output\_rows &= input\_rows - kernel\_rows + 1 \\ output\_cols &= input\_cols - kernel\_cols + 1 \end{aligned}$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 中。

示例

示例 1

Input:  体积 V∈R^{3×3×3}:
          V(d=0) = [[1,2,3],[4,5,6],[7,8,9]]
          V(d=1) = [[10,11,12],[13,14,15],[16,17,18]]
          V(d=2) = [[19,20,21],[22,23,24],[25,26,27]]
        Kernel K∈R^{2×3×3}:
          K(d=0) = [[1,0,0],[1,1,1],[0,0,0]]
          K(d=1) = [[1,1,0],[1,1,0],[0,0,1]]
Output: O∈R^{2×1×1}: [82, 163]

示例 2

Input:  体积 V∈R^{2×2×2}:
          V(d=0) = [[1,2],[3,4]]
          V(d=1) = [[5,6],[7,8]]
        Kernel K∈R^{2×2×2} (全1)
Output: O∈R^{1×1×1}: [36]

约束条件

• $1 \le input\_depth, input\_rows, input\_cols \le 256$。
• $1 \le kernel\_depth, kernel\_rows, kernel\_cols \le 5$。
• 卷积核各维度不超过输入对应维度。
• 性能测试在 $input = 128 \times 128$，$kernel = 5 \times 5$ 的规模下进行。

解题思路

三维卷积是二维卷积的自然扩展，常用于视频处理（3D CNN）和医学影像。计算量随核体积呈立方增长。共享内存分块策略在三维空间中更加关键——因为 $K_d \times K_r \times K_c$ 个输入元素可能超过每个 SM 的共享内存限制，需要仔细规划分块大小。常量内存对小核（$\le 5^3 = 125$ 个参数）仍然有效。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void conv3d_kernel(const float* input, const float* kernel, float* output,
    int D, int H, int W, int KD, int KH, int KW) {
    int oc = blockIdx.x * blockDim.x + threadIdx.x;
    int or_ = blockIdx.y * blockDim.y + threadIdx.y;
    int od = blockIdx.z * blockDim.z + threadIdx.z;
    int OD = D-KD+1, OH = H-KH+1, OW = W-KW+1;
    if (od < OD && or_ < OH && oc < OW) {
        float sum = 0.0f;
        for (int kd = 0; kd < KD; kd++)
            for (int ky = 0; ky < KH; ky++)
                for (int kx = 0; kx < KW; kx++)
                    sum += input[((od+kd)*H + (or_+ky))*W + (oc+kx)]
                         * kernel[(kd*KH + ky)*KW + kx];
        output[(od*OH + or_)*OW + oc] = sum;
    }
}
extern "C" void solve(const float* input, const float* kernel, float* output,
    int D, int H, int W, int KD, int KH, int KW) {
    int OD=D-KD+1, OH=H-KH+1, OW=W-KW+1;
    dim3 threads(8,8,4), blocks((OW+7)/8, (OH+7)/8, (OD+3)/4);
    conv3d_kernel<<<blocks, threads>>>(input, kernel, output, D, H, W, KD, KH, KW);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl

@triton.jit
def conv3d_kernel(input_ptr, kernel_ptr, output_ptr, D, H, W, KD, KH, KW, BLOCK: tl.constexpr):
    ow = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    oh = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
    od = tl.program_id(2)
    OW = W - KW + 1
    OH = H - KH + 1
    OD_ = D - KD + 1
    mask = (od < OD_) & (oh[:, None] < OH) & (ow[None, :] < OW)
    acc = tl.zeros((BLOCK, BLOCK), tl.float32)
    for kd in range(KD):
        for ky in range(KH):
            for kx in range(KW):
                val = tl.load(input_ptr + ((od + kd) * H + oh[:, None] + ky) * W + ow[None, :] + kx, mask=mask)
                kw = tl.load(kernel_ptr + (kd * KH + ky) * KW + kx)
                acc += val * kw
    tl.store(output_ptr + (od * OH + oh[:, None]) * OW + ow[None, :], acc, mask=mask)