forked from ccf-ai-infra/GPUCodeForces
213 lines
7.0 KiB
Python
213 lines
7.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
BATCH_SIZE = 8
|
|
CHANNELS = 16
|
|
DEPTH = 16
|
|
HEIGHT = 16
|
|
WIDTH = 16
|
|
|
|
PADDING = (1, 1, 2, 2, 1, 0)
|
|
|
|
BLOCK_DIM_X = 8
|
|
BLOCK_DIM_Y = 8
|
|
BLOCK_DIM_Z = 8
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = (padding,) * 6
|
|
else:
|
|
self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = padding
|
|
|
|
self.block_dim_x = BLOCK_DIM_X
|
|
self.block_dim_y = BLOCK_DIM_Y
|
|
self.block_dim_z = BLOCK_DIM_Z
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor reflection_pad3d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B,
|
|
int pad_F, int pad_K
|
|
);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
#define BLOCK_DIM_X {self.block_dim_x}
|
|
#define BLOCK_DIM_Y {self.block_dim_y}
|
|
#define BLOCK_DIM_Z {self.block_dim_z}
|
|
|
|
|
|
__device__ inline int reflect_idx(
|
|
int j, int pad_before, int W_in
|
|
) {{
|
|
if (j < pad_before) {{
|
|
return pad_before - j;
|
|
}} else if (j < (pad_before + W_in)) {{
|
|
return j - pad_before;
|
|
}} else {{
|
|
int j_rel = j - (pad_before + W_in);
|
|
return W_in - 2 - j_rel;
|
|
}}
|
|
}}
|
|
|
|
|
|
__global__ void reflection_pad3d_fused_kernel(
|
|
const float* __restrict__ input_data,
|
|
float* __restrict__ output_data,
|
|
int N, int C,
|
|
int D_in, int H_in, int W_in,
|
|
int D_out, int H_out, int W_out,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B,
|
|
int pad_F, int pad_K
|
|
) {{
|
|
extern __shared__ float s_in[];
|
|
|
|
const int n_idx = blockIdx.x;
|
|
const int c_idx = blockIdx.y;
|
|
const int tid_x = threadIdx.x;
|
|
const int tid_y = threadIdx.y;
|
|
const int tid_z = threadIdx.z;
|
|
|
|
|
|
const int64_t H_in_stride = W_in;
|
|
const int64_t D_in_stride = H_in * W_in;
|
|
|
|
const int64_t C_in_stride = D_in * D_in_stride;
|
|
|
|
|
|
const int64_t H_out_stride = W_out;
|
|
const int64_t D_out_stride = H_out * W_out;
|
|
|
|
const int64_t C_out_stride = D_out * D_out_stride;
|
|
|
|
|
|
const float* p_in_base = input_data + (n_idx * C + c_idx) * C_in_stride;
|
|
float* p_out_base = output_data + (n_idx * C + c_idx) * C_out_stride;
|
|
|
|
|
|
for (int k = tid_z; k < D_in; k += BLOCK_DIM_Z) {{
|
|
for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
|
|
for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
|
|
// (k, i, j) -> 1D index
|
|
int64_t in_idx = k*D_in_stride + i*H_in_stride + j;
|
|
s_in[in_idx] = p_in_base[in_idx];
|
|
}}
|
|
}}
|
|
}}
|
|
__syncthreads(); // 确保 s_in 加载完成
|
|
|
|
|
|
for (int k = tid_z; k < D_out; k += BLOCK_DIM_Z) {{
|
|
int in_k = reflect_idx(k, pad_F, D_in);
|
|
|
|
for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{
|
|
int in_i = reflect_idx(i, pad_T, H_in);
|
|
|
|
for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{
|
|
int in_j = reflect_idx(j, pad_L, W_in);
|
|
|
|
// 从共享内存读取 (使用 IN 步长)
|
|
int64_t s_in_idx = in_k*D_in_stride + in_i*H_in_stride + in_j;
|
|
|
|
// 写入全局内存 (使用 OUT 步长)
|
|
int64_t p_out_idx = k*D_out_stride + i*H_out_stride + j;
|
|
|
|
p_out_base[p_out_idx] = s_in[s_in_idx];
|
|
}}
|
|
}}
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor reflection_pad3d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B,
|
|
int pad_F, int pad_K
|
|
) {{
|
|
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
|
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
|
TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)");
|
|
|
|
const int64_t N_64 = input.size(0);
|
|
const int64_t C_64 = input.size(1);
|
|
const int64_t D_in_64 = input.size(2);
|
|
const int64_t H_in_64 = input.size(3);
|
|
const int64_t W_in_64 = input.size(4);
|
|
|
|
TORCH_CHECK(pad_L < W_in_64, "pad_L error");
|
|
TORCH_CHECK(pad_R < W_in_64, "pad_R error");
|
|
TORCH_CHECK(pad_T < H_in_64, "pad_T error");
|
|
TORCH_CHECK(pad_B < H_in_64, "pad_B error");
|
|
TORCH_CHECK(pad_F < D_in_64, "pad_F error");
|
|
TORCH_CHECK(pad_K < D_in_64, "pad_K error");
|
|
|
|
const int64_t D_out_64 = D_in_64 + pad_F + pad_K;
|
|
const int64_t H_out_64 = H_in_64 + pad_T + pad_B;
|
|
const int64_t W_out_64 = W_in_64 + pad_L + pad_R;
|
|
|
|
auto output = torch::empty({{N_64, C_64, D_out_64, H_out_64, W_out_64}}, input.options());
|
|
|
|
dim3 grid_dim(N_64, C_64);
|
|
dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z);
|
|
|
|
|
|
const int shared_mem_size = D_in_64 * H_in_64 * W_in_64 * sizeof(float);
|
|
|
|
reflection_pad3d_fused_kernel<<<grid_dim, block_dim, shared_mem_size>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
static_cast<int>(N_64), static_cast<int>(C_64),
|
|
static_cast<int>(D_in_64), static_cast<int>(H_in_64), static_cast<int>(W_in_64),
|
|
static_cast<int>(D_out_64), static_cast<int>(H_out_64), static_cast<int>(W_out_64),
|
|
pad_L, pad_R,
|
|
pad_T, pad_B,
|
|
pad_F, pad_K
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = [
|
|
'-O3',
|
|
'--use_fast_math',
|
|
'--expt-relaxed-constexpr'
|
|
]
|
|
|
|
self.pad_op = load_inline(
|
|
name="reflection_pad3d_op_v2_fixed",
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["reflection_pad3d_forward_cuda"],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
x_cont = x.contiguous()
|
|
|
|
return self.pad_op.reflection_pad3d_forward_cuda(
|
|
x_cont,
|
|
self.pad_L, self.pad_R,
|
|
self.pad_T, self.pad_B,
|
|
self.pad_F, self.pad_K
|
|
) |