forked from ccf-ai-infra/GPUCodeForces
373 lines
12 KiB
Python
373 lines
12 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义 (与你之前的代码一致)
|
|
# -------------------------------------------------------------
|
|
BATCH_SIZE = 32
|
|
CHANNELS = 64
|
|
HEIGHT = 32 # H_in
|
|
WIDTH = 32 # W_in
|
|
PADDING = (1, 1, 2, 0)
|
|
BLOCK_DIM_X = 16
|
|
BLOCK_DIM_Y = 16
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class ModelNew(nn.Module):
|
|
"""
|
|
ReflectionPad2d 的高性能 CUDA 融合核函数实现
|
|
(V3: 修复了 'contiguous' 运行时错误)
|
|
"""
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
self.pad_L = padding
|
|
self.pad_R = padding
|
|
self.pad_T = padding
|
|
self.pad_B = padding
|
|
else:
|
|
self.pad_L = padding[0]
|
|
self.pad_R = padding[1]
|
|
self.pad_T = padding[2]
|
|
self.pad_B = padding[3]
|
|
|
|
self.block_dim_x = BLOCK_DIM_X
|
|
self.block_dim_y = BLOCK_DIM_Y
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor reflection_pad2d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
);
|
|
"""
|
|
|
|
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}
|
|
|
|
__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_pad2d_fused_kernel(
|
|
const float* __restrict__ input_data,
|
|
float* __restrict__ output_data,
|
|
int N, int C,
|
|
int H_in, int W_in,
|
|
int H_out, int W_out,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
) {{
|
|
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 float* p_in = input_data + (n_idx * C + c_idx) * (H_in * W_in);
|
|
float* p_out = output_data + (n_idx * C + c_idx) * (H_out * W_out);
|
|
|
|
// Pass 1: Load to shared memory
|
|
for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
|
|
for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
|
|
s_in[i * W_in + j] = p_in[i * W_in + j];
|
|
}}
|
|
}}
|
|
__syncthreads();
|
|
|
|
// Pass 2: Compute and store from shared memory
|
|
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);
|
|
p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
|
|
}}
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor reflection_pad2d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
) {{
|
|
// 这个检查现在是安全的，因为我们在 Python 中确保了连续性
|
|
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
|
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
|
TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
|
|
|
|
const int64_t N_64 = input.size(0);
|
|
const int64_t C_64 = input.size(1);
|
|
const int64_t H_in_64 = input.size(2);
|
|
const int64_t W_in_64 = input.size(3);
|
|
|
|
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");
|
|
|
|
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, H_out_64, W_out_64}}, input.options());
|
|
|
|
dim3 grid_dim(N_64, C_64);
|
|
dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);
|
|
|
|
const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);
|
|
|
|
reflection_pad2d_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>(H_in_64), static_cast<int>(W_in_64),
|
|
static_cast<int>(H_out_64), static_cast<int>(W_out_64),
|
|
pad_L, pad_R,
|
|
pad_T, pad_B
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = [
|
|
'-O3',
|
|
'--use_fast_math',
|
|
'--expt-relaxed-constexpr'
|
|
]
|
|
|
|
self.pad_op = load_inline(
|
|
name="reflection_pad2d_op_v2_fixed", # (与 V2 编译的二进制文件相同)
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["reflection_pad2d_forward_cuda"],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
# [修复]
|
|
# 必须确保张量是连续的，才能传递给 C++/CUDA
|
|
x_cont = x.contiguous()
|
|
|
|
# 调用我们编译好的 CUDA C++ 函数
|
|
return self.pad_op.reflection_pad2d_forward_cuda(
|
|
x_cont, # 传递连续的张量
|
|
self.pad_L, self.pad_R,
|
|
self.pad_T, self.pad_B
|
|
)
|
|
import torch
|
|
|
|
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义 (与你之前的代码一致)
|
|
# -------------------------------------------------------------
|
|
BATCH_SIZE = 32
|
|
CHANNELS = 64
|
|
HEIGHT = 32 # H_in
|
|
WIDTH = 32 # W_in
|
|
PADDING = (1, 1, 2, 0)
|
|
BLOCK_DIM_X = 16
|
|
BLOCK_DIM_Y = 16
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class ModelNew(nn.Module):
|
|
"""
|
|
ReflectionPad2d 的高性能 CUDA 融合核函数实现
|
|
(V3: 修复了 'contiguous' 运行时错误)
|
|
"""
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
self.pad_L = padding
|
|
self.pad_R = padding
|
|
self.pad_T = padding
|
|
self.pad_B = padding
|
|
else:
|
|
self.pad_L = padding[0]
|
|
self.pad_R = padding[1]
|
|
self.pad_T = padding[2]
|
|
self.pad_B = padding[3]
|
|
|
|
self.block_dim_x = BLOCK_DIM_X
|
|
self.block_dim_y = BLOCK_DIM_Y
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor reflection_pad2d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
);
|
|
"""
|
|
|
|
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}
|
|
|
|
__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_pad2d_fused_kernel(
|
|
const float* __restrict__ input_data,
|
|
float* __restrict__ output_data,
|
|
int N, int C,
|
|
int H_in, int W_in,
|
|
int H_out, int W_out,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
) {{
|
|
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 float* p_in = input_data + (n_idx * C + c_idx) * (H_in * W_in);
|
|
float* p_out = output_data + (n_idx * C + c_idx) * (H_out * W_out);
|
|
|
|
// Pass 1: Load to shared memory
|
|
for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
|
|
for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
|
|
s_in[i * W_in + j] = p_in[i * W_in + j];
|
|
}}
|
|
}}
|
|
__syncthreads();
|
|
|
|
// Pass 2: Compute and store from shared memory
|
|
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);
|
|
p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
|
|
}}
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor reflection_pad2d_forward_cuda(
|
|
torch::Tensor input,
|
|
int pad_L, int pad_R,
|
|
int pad_T, int pad_B
|
|
) {{
|
|
// 这个检查现在是安全的，因为我们在 Python 中确保了连续性
|
|
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
|
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
|
TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
|
|
|
|
const int64_t N_64 = input.size(0);
|
|
const int64_t C_64 = input.size(1);
|
|
const int64_t H_in_64 = input.size(2);
|
|
const int64_t W_in_64 = input.size(3);
|
|
|
|
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");
|
|
|
|
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, H_out_64, W_out_64}}, input.options());
|
|
|
|
dim3 grid_dim(N_64, C_64);
|
|
dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);
|
|
|
|
const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);
|
|
|
|
reflection_pad2d_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>(H_in_64), static_cast<int>(W_in_64),
|
|
static_cast<int>(H_out_64), static_cast<int>(W_out_64),
|
|
pad_L, pad_R,
|
|
pad_T, pad_B
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = [
|
|
'-O3',
|
|
'--use_fast_math',
|
|
'--expt-relaxed-constexpr'
|
|
]
|
|
|
|
self.pad_op = load_inline(
|
|
name="reflection_pad2d_op_v2_fixed", # (与 V2 编译的二进制文件相同)
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["reflection_pad2d_forward_cuda"],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
# [修复]
|
|
# 必须确保张量是连续的，才能传递给 C++/CUDA
|
|
x_cont = x.contiguous()
|
|
|
|
# 调用我们编译好的 CUDA C++ 函数
|
|
return self.pad_op.reflection_pad2d_forward_cuda(
|
|
x_cont, # 传递连续的张量
|
|
self.pad_L, self.pad_R,
|
|
self.pad_T, self.pad_B
|
|
)
|