forked from ccf-ai-infra/GPUCodeForces
145 lines
4.6 KiB
Python
145 lines
4.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
BATCH_SIZE = 32
|
|
CHANNELS = 64
|
|
WIDTH = 128 # W_in
|
|
PADDING = (3, 1) # (padding_left, padding_right)
|
|
BLOCK_SIZE = 256 # CUDA Block 维度
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class ModelNew(nn.Module):
|
|
"""
|
|
ReflectionPad1d 的高性能 CUDA 融合核函数实现
|
|
(修复了编译错误)
|
|
"""
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
self.pad_L = padding
|
|
self.pad_R = padding
|
|
else:
|
|
self.pad_L = padding[0]
|
|
self.pad_R = padding[1]
|
|
|
|
self.block_size = BLOCK_SIZE
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor reflection_pad1d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L,
|
|
int pad_R
|
|
);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
// [修复] 将 #define 移至此处
|
|
#define BLOCK_SIZE {self.block_size}
|
|
|
|
/*
|
|
* ReflectionPad1d 融合核函数
|
|
*/
|
|
__global__ void reflection_pad1d_fused_kernel(
|
|
const float* __restrict__ input_data,
|
|
float* __restrict__ output_data,
|
|
int N, int C, int W_in, int W_out,
|
|
int pad_L, int pad_R
|
|
) {{ // <-- f-string 转义
|
|
const int n_idx = blockIdx.x;
|
|
const int c_idx = blockIdx.y;
|
|
const int tid = threadIdx.x;
|
|
|
|
const float* p_in = input_data + (n_idx * C + c_idx) * W_in;
|
|
float* p_out = output_data + (n_idx * C + c_idx) * W_out;
|
|
|
|
// [修复] BLOCK_SIZE 现在可见
|
|
for (int j = tid; j < W_out; j += BLOCK_SIZE) {{ // <-- f-string 转义
|
|
int in_idx = 0;
|
|
|
|
if (j < pad_L) {{
|
|
in_idx = pad_L - j;
|
|
}} else if (j < (pad_L + W_in)) {{
|
|
in_idx = j - pad_L;
|
|
}} else {{
|
|
int j_rel = j - (pad_L + W_in);
|
|
in_idx = W_in - 2 - j_rel;
|
|
}}
|
|
|
|
p_out[j] = p_in[in_idx];
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
// [修复] torch.Tensor -> torch::Tensor
|
|
torch::Tensor reflection_pad1d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L,
|
|
int pad_R
|
|
) {{
|
|
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
|
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
|
TORCH_CHECK(input.dim() == 3, "input must be 3D (N, C, W)");
|
|
|
|
const int64_t N_64 = input.size(0);
|
|
const int64_t C_64 = input.size(1);
|
|
const int64_t W_in_64 = input.size(2);
|
|
|
|
TORCH_CHECK(pad_L < W_in_64, "padding_left should be less than input width");
|
|
TORCH_CHECK(pad_R < W_in_64, "padding_right should be less than input width");
|
|
|
|
const int64_t W_out_64 = W_in_64 + pad_L + pad_R;
|
|
|
|
auto output = torch::empty({{N_64, C_64, W_out_64}}, input.options());
|
|
|
|
dim3 grid_dim(N_64, C_64);
|
|
dim3 block_dim(BLOCK_SIZE);
|
|
|
|
reflection_pad1d_fused_kernel<<<grid_dim, block_dim>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
static_cast<int>(N_64),
|
|
static_cast<int>(C_64),
|
|
static_cast<int>(W_in_64),
|
|
static_cast<int>(W_out_64),
|
|
pad_L,
|
|
pad_R
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
# JIT (Just-In-Time) 编译
|
|
self.pad_op = load_inline(
|
|
name="reflection_pad1d_op_v3_fixed", # 更改名称以避免缓存
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["reflection_pad1d_forward_cuda"],
|
|
verbose=False # 如果还报错，请设为 True
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
# 调用我们编译好的 CUDA C++ 函数
|
|
return self.pad_op.reflection_pad1d_forward_cuda(
|
|
x,
|
|
self.pad_L,
|
|
self.pad_R
|
|
) |