forked from ccf-ai-infra/GPUCodeForces
305 lines
9.6 KiB
Python
305 lines
9.6 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
instancenorm_source = """
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
#include <math.h>
|
||
|
||
const int WARP_SIZE = 32;
|
||
|
||
// 优化的warp级归约
|
||
__inline__ __device__ float warpReduceSum(float val) {
|
||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||
}
|
||
return val;
|
||
}
|
||
|
||
__inline__ __device__ float warpReduceSumSq(float val) {
|
||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||
}
|
||
return val;
|
||
}
|
||
|
||
// 优化的InstanceNorm内核 - 使用warp级和block级混合归约
|
||
__global__ void instancenorm_optimized_kernel(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ weight,
|
||
const float* __restrict__ bias,
|
||
float* __restrict__ y,
|
||
int batch,
|
||
int channels,
|
||
int height,
|
||
int width,
|
||
float eps
|
||
) {
|
||
int spatial_size = height * width;
|
||
int instance_idx = blockIdx.x;
|
||
int channel_idx = blockIdx.y;
|
||
|
||
if (instance_idx >= batch || channel_idx >= channels) return;
|
||
|
||
int tid = threadIdx.x;
|
||
int lane_id = tid % WARP_SIZE;
|
||
int warp_id = tid / WARP_SIZE;
|
||
int instance_offset = instance_idx * channels * spatial_size + channel_idx * spatial_size;
|
||
|
||
extern __shared__ float sdata[];
|
||
float* warp_sums = sdata;
|
||
float* warp_sum_sqs = sdata + (blockDim.x / WARP_SIZE) * 2;
|
||
|
||
// 第一阶段：每个warp内部归约
|
||
float sum = 0.0f;
|
||
float sum_sq = 0.0f;
|
||
|
||
// 使用循环展开和向量化友好的访问模式
|
||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||
float v = x[instance_offset + i];
|
||
sum += v;
|
||
sum_sq += v * v;
|
||
}
|
||
|
||
// Warp级归约
|
||
sum = warpReduceSum(sum);
|
||
sum_sq = warpReduceSumSq(sum_sq);
|
||
|
||
// 每个warp的第一个线程保存结果到shared memory
|
||
if (lane_id == 0) {
|
||
warp_sums[warp_id] = sum;
|
||
warp_sum_sqs[warp_id] = sum_sq;
|
||
}
|
||
__syncthreads();
|
||
|
||
// 第二阶段：block级归约（只在warp 0中进行）
|
||
if (warp_id == 0) {
|
||
sum = (lane_id < (blockDim.x / WARP_SIZE)) ? warp_sums[lane_id] : 0.0f;
|
||
sum_sq = (lane_id < (blockDim.x / WARP_SIZE)) ? warp_sum_sqs[lane_id] : 0.0f;
|
||
|
||
// 再次warp归约
|
||
sum = warpReduceSum(sum);
|
||
sum_sq = warpReduceSumSq(sum_sq);
|
||
|
||
// 计算最终统计量
|
||
if (lane_id == 0) {
|
||
float mean = sum / spatial_size;
|
||
float var = (sum_sq / spatial_size) - (mean * mean);
|
||
var = fmaxf(var, 0.0f);
|
||
|
||
// 保存到shared memory供所有线程使用
|
||
warp_sums[0] = mean;
|
||
warp_sum_sqs[0] = rsqrtf(var + eps);
|
||
warp_sums[1] = weight[channel_idx];
|
||
warp_sum_sqs[1] = bias[channel_idx];
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
// 所有线程读取统计量
|
||
float mean = warp_sums[0];
|
||
float inv_std = warp_sum_sqs[0];
|
||
float w = warp_sums[1];
|
||
float b = warp_sum_sqs[1];
|
||
|
||
// 应用InstanceNorm - 使用更优化的内存访问模式
|
||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||
float v = x[instance_offset + i];
|
||
float norm_val = (v - mean) * inv_std;
|
||
y[instance_offset + i] = norm_val * w + b;
|
||
}
|
||
}
|
||
|
||
// 针对小尺寸的优化内核
|
||
__global__ void instancenorm_small_kernel(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ weight,
|
||
const float* __restrict__ bias,
|
||
float* __restrict__ y,
|
||
int batch,
|
||
int channels,
|
||
int height,
|
||
int width,
|
||
float eps
|
||
) {
|
||
int spatial_size = height * width;
|
||
int instance_idx = blockIdx.x;
|
||
int channel_idx = blockIdx.y;
|
||
|
||
if (instance_idx >= batch || channel_idx >= channels) return;
|
||
|
||
int tid = threadIdx.x;
|
||
int instance_offset = instance_idx * channels * spatial_size + channel_idx * spatial_size;
|
||
|
||
extern __shared__ float sdata[];
|
||
float* sum_shared = sdata;
|
||
float* sum_sq_shared = sdata + blockDim.x;
|
||
|
||
// 针对小尺寸的简化归约
|
||
float sum = 0.0f;
|
||
float sum_sq = 0.0f;
|
||
|
||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||
float v = x[instance_offset + i];
|
||
sum += v;
|
||
sum_sq += v * v;
|
||
}
|
||
|
||
sum_shared[tid] = sum;
|
||
sum_sq_shared[tid] = sum_sq;
|
||
__syncthreads();
|
||
|
||
// 树状归约
|
||
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
||
if (tid < offset) {
|
||
sum_shared[tid] += sum_shared[tid + offset];
|
||
sum_sq_shared[tid] += sum_sq_shared[tid + offset];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
__shared__ float s_mean;
|
||
__shared__ float s_inv_std;
|
||
__shared__ float s_weight;
|
||
__shared__ float s_bias;
|
||
|
||
if (tid == 0) {
|
||
float mean = sum_shared[0] / spatial_size;
|
||
float var = (sum_sq_shared[0] / spatial_size) - (mean * mean);
|
||
var = fmaxf(var, 0.0f);
|
||
s_mean = mean;
|
||
s_inv_std = rsqrtf(var + eps);
|
||
s_weight = weight[channel_idx];
|
||
s_bias = bias[channel_idx];
|
||
}
|
||
__syncthreads();
|
||
|
||
float mean = s_mean;
|
||
float inv_std = s_inv_std;
|
||
float w = s_weight;
|
||
float b = s_bias;
|
||
|
||
// 应用归一化
|
||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||
float v = x[instance_offset + i];
|
||
float norm_val = (v - mean) * inv_std;
|
||
y[instance_offset + i] = norm_val * w + b;
|
||
}
|
||
}
|
||
|
||
torch::Tensor instancenorm_cuda_forward(
|
||
torch::Tensor x,
|
||
torch::Tensor weight,
|
||
torch::Tensor bias,
|
||
float eps
|
||
) {
|
||
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
|
||
TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
|
||
TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
|
||
TORCH_CHECK(x.dim() == 4, "input must be 4D [batch, channels, height, width]");
|
||
TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
|
||
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
|
||
TORCH_CHECK(x.size(1) == weight.size(0), "channel size mismatch");
|
||
TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias size mismatch");
|
||
|
||
auto x_contig = x.contiguous();
|
||
int batch = x_contig.size(0);
|
||
int channels = x_contig.size(1);
|
||
int height = x_contig.size(2);
|
||
int width = x_contig.size(3);
|
||
int spatial_size = height * width;
|
||
|
||
auto y = torch::empty_like(x_contig);
|
||
|
||
// 更智能的线程配置
|
||
dim3 blocks(batch, channels);
|
||
size_t shared_mem;
|
||
|
||
if (spatial_size >= 1024) {
|
||
// 大尺寸使用优化内核，256线程，4个warp
|
||
int threads = 256;
|
||
shared_mem = (threads / WARP_SIZE) * 2 * sizeof(float) + 4 * sizeof(float);
|
||
instancenorm_optimized_kernel<<<blocks, threads, shared_mem>>>(
|
||
x_contig.data_ptr<float>(),
|
||
weight.data_ptr<float>(),
|
||
bias.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
batch, channels, height, width, eps
|
||
);
|
||
} else {
|
||
// 小尺寸使用简化内核
|
||
int threads;
|
||
if (spatial_size <= 64) threads = 64;
|
||
else if (spatial_size <= 128) threads = 128;
|
||
else threads = 256;
|
||
|
||
threads = min(threads, spatial_size);
|
||
if (threads < 32) threads = 32;
|
||
|
||
shared_mem = 2 * threads * sizeof(float);
|
||
instancenorm_small_kernel<<<blocks, threads, shared_mem>>>(
|
||
x_contig.data_ptr<float>(),
|
||
weight.data_ptr<float>(),
|
||
bias.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
batch, channels, height, width, eps
|
||
);
|
||
}
|
||
|
||
// 移除同步，让CUDA流自动管理
|
||
// cudaDeviceSynchronize();
|
||
return y;
|
||
}
|
||
"""
|
||
|
||
instancenorm_cpp_source = """
|
||
torch::Tensor instancenorm_cuda_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps);
|
||
"""
|
||
|
||
instancenorm_cuda = load_inline(
|
||
name="instancenorm_cuda",
|
||
cpp_sources=instancenorm_cpp_source,
|
||
cuda_sources=instancenorm_source,
|
||
functions=["instancenorm_cuda_forward"],
|
||
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
||
verbose=True
|
||
)
|
||
|
||
|
||
# CUDA优化版本 - 完全与PyTorch一致
|
||
class ModelNew(nn.Module):
|
||
"""
|
||
Simplified CUDA version that forces track_running_stats=False for exact equivalence.
|
||
"""
|
||
|
||
def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
|
||
super(ModelNew, self).__init__()
|
||
|
||
# 强制track_running_stats=False以确保与CUDA实现完全等价
|
||
if track_running_stats:
|
||
print("警告：CUDA优化版本不支持track_running_stats=True，已强制设置为False")
|
||
|
||
self.num_features = num_features
|
||
self.eps = eps
|
||
self.affine = affine
|
||
self.track_running_stats = False # 强制为False
|
||
|
||
if affine:
|
||
self.weight = nn.Parameter(torch.ones(num_features))
|
||
self.bias = nn.Parameter(torch.zeros(num_features))
|
||
else:
|
||
self.register_parameter('weight', None)
|
||
self.register_parameter('bias', None)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
只支持track_running_stats=False的情况
|
||
"""
|
||
if self.affine:
|
||
return instancenorm_cuda.instancenorm_cuda_forward(x, self.weight, self.bias, self.eps)
|
||
else:
|
||
weight = torch.ones(self.num_features, device=x.device)
|
||
bias = torch.zeros(self.num_features, device=x.device)
|
||
return instancenorm_cuda.instancenorm_cuda_forward(x, weight, bias, self.eps) |