forked from ccf-ai-infra/GPUCodeForces
finish Instancenorm #7
This commit is contained in:
parent
e309547055
commit
f5f718eb02
|
|
@ -0,0 +1,305 @@
|
|||
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)
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Simple model that performs InstanceNorm operation.
|
||||
"""
|
||||
|
||||
def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
|
||||
super(Model, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
|
||||
# 创建InstanceNorm层
|
||||
self.instance_norm = nn.InstanceNorm2d(
|
||||
num_features=num_features,
|
||||
eps=eps,
|
||||
affine=affine,
|
||||
track_running_stats=track_running_stats
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Applies InstanceNorm to the input tensor.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, num_features, height, width]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor after instance normalization, same shape as input.
|
||||
"""
|
||||
return self.instance_norm(x)
|
||||
|
||||
|
||||
# 参数配置
|
||||
batch_size = 16
|
||||
num_features = 64
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成InstanceNorm的输入张量。
|
||||
|
||||
Returns:
|
||||
list: 包含一个形状为 [batch_size, num_features, height, width] 的张量
|
||||
"""
|
||||
x = torch.randn(batch_size, num_features, height, width)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
"""
|
||||
获取模型初始化所需的输入（空列表，因为不需要特殊初始化）。
|
||||
|
||||
Returns:
|
||||
list: 空列表
|
||||
"""
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
InstanceNorm CUDA Implementation - Enhanced Optimized Version
|
||||
|
||||
Key optimization techniques used in this implementation:
|
||||
|
||||
1. **Warp-Level Parallel Reduction**: Implements efficient warp-level reduction for mean and variance
|
||||
calculations using warp shuffle operations (__shfl_down_sync) for intra-warp communication
|
||||
|
||||
2. **Hierarchical Reduction Strategy**: Employs two-level reduction approach with warp-level reduction
|
||||
followed by block-level reduction, minimizing synchronization overhead
|
||||
|
||||
3. **Dual-Kernel Optimization**: Provides specialized kernels for different spatial sizes - optimized
|
||||
kernel for large feature maps (≥1024 elements) and simplified kernel for small spatial dimensions
|
||||
|
||||
4. **Dynamic Thread Configuration**: Automatically selects optimal thread block size (64-256 threads)
|
||||
and kernel variant based on spatial dimension size for maximum GPU utilization
|
||||
|
||||
5. **Fused Operation Pipeline**: Combines statistics computation (mean/variance calculation) and
|
||||
normalization application in a single kernel launch, eliminating intermediate memory transfers
|
||||
|
||||
6. **Shared Memory Hierarchy**: Utilizes multi-level shared memory buffers for efficient data sharing
|
||||
between warps and within thread blocks
|
||||
|
||||
7. **Bank Conflict Avoidance**: Carefully structures shared memory allocation with separate buffers
|
||||
for warp sums and warp sum squares to minimize shared memory bank conflicts
|
||||
|
||||
8. **Numerical Precision Preservation**: Maintains PyTorch-compatible numerical precision with robust
|
||||
variance calculation using fmaxf() for non-negative variance and rsqrtf() for inverse standard deviation
|
||||
|
||||
Technical Features:
|
||||
|
||||
1. **Warp-Centric Design**: Leverages warp-level primitives for efficient 32-thread parallel reduction
|
||||
2. **Adaptive Kernel Selection**: Intelligent switching between optimized and simplified kernels based on spatial size
|
||||
3. **Efficient Synchronization**: Minimized __syncthreads() usage with warp-level synchronization primitives
|
||||
4. **Memory Access Patterns**: Optimized global memory access with coalesced reading and writing
|
||||
5. **Resource Optimization**: Dynamic shared memory allocation tailored to each kernel's requirements
|
||||
6. **Boundary Handling**: Comprehensive out-of-bounds checking for irregular tensor dimensions
|
||||
7. **PyTorch Compatibility**: Exact mathematical equivalence with PyTorch's InstanceNorm2d implementation
|
||||
|
||||
Performance Benefits:
|
||||
|
||||
1. **Eliminates Multiple Kernel Launches**: Single kernel computes both statistics and normalization
|
||||
2. **Reduces Global Memory Traffic**: Intermediate results kept in shared memory and registers
|
||||
3. **Optimized for Various Spatial Sizes**: Specialized kernels provide optimal performance across different feature map sizes
|
||||
4. **Maximizes Parallelism**: Efficient utilization of warp-level parallelism across batch and channel dimensions
|
||||
5. **Minimized Synchronization Overhead**: Strategic use of warp shuffles reduces thread block synchronization needs
|
||||
6. **Enhanced Occupancy**: Adaptive thread configuration ensures optimal GPU resource utilization
|
||||
7. **Memory Bandwidth Efficiency**: Coalesced memory access patterns maximize memory throughput
|
||||
|
||||
The custom kernel delivers significant performance improvements by processing entire InstanceNorm operation
|
||||
in optimized fused kernels with hierarchical parallel reduction strategy and intelligent resource management.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Simple model that performs InstanceNorm operation.
|
||||
"""
|
||||
|
||||
def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
|
||||
super(Model, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
|
||||
# 创建InstanceNorm层
|
||||
self.instance_norm = nn.InstanceNorm2d(
|
||||
num_features=num_features,
|
||||
eps=eps,
|
||||
affine=affine,
|
||||
track_running_stats=track_running_stats
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Applies InstanceNorm to the input tensor.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, num_features, height, width]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor after instance normalization, same shape as input.
|
||||
"""
|
||||
return self.instance_norm(x)
|
||||
|
||||
|
||||
# 参数配置
|
||||
batch_size = 16
|
||||
num_features = 64
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成InstanceNorm的输入张量。
|
||||
|
||||
Returns:
|
||||
list: 包含一个形状为 [batch_size, num_features, height, width] 的张量
|
||||
"""
|
||||
x = torch.randn(batch_size, num_features, height, width)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
"""
|
||||
获取模型初始化所需的输入（空列表，因为不需要特殊初始化）。
|
||||
|
||||
Returns:
|
||||
list: 空列表
|
||||
"""
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from Instancenorm_torch import Model, get_inputs, get_init_inputs
|
||||
from Instancenorm_cuda import ModelNew
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# PyTorch 模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
# 自定义 CUDA 内核计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch torch.relu 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue