forked from ccf-ai-infra/GPUCodeForces
feat:add minkowski_instancenorm #78
This commit is contained in:
parent
f989885dde
commit
e57cb5cc4d
|
|
@ -0,0 +1,132 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
instancenorm_minkowski_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
// 线程块内归约求和
|
||||
__inline__ __device__ float blockReduceSum(float val, int tid) {
|
||||
extern __shared__ float sdata[];
|
||||
sdata[tid] = val;
|
||||
__syncthreads();
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
sdata[tid] += sdata[tid + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
return sdata[0];
|
||||
}
|
||||
|
||||
// 融合内核 - InstanceNorm -> Minkowski
|
||||
__global__ void instancenorm_minkowski_kernel(
|
||||
const float* __restrict__ x,
|
||||
float* __restrict__ y,
|
||||
int N, int C, int H, int W,
|
||||
float p, float inv_p, float eps
|
||||
) {
|
||||
// blockIdx.x: sample index (N)
|
||||
// blockIdx.y: channel index (C)
|
||||
int sample_idx = blockIdx.x;
|
||||
int channel_idx = blockIdx.y;
|
||||
|
||||
if (sample_idx >= N || channel_idx >= C) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int spatial_size = H * W;
|
||||
|
||||
// 计算当前样本和通道的起始偏移量
|
||||
int base_offset = sample_idx * C * spatial_size + channel_idx * spatial_size;
|
||||
|
||||
// --- 融合步骤1：计算InstanceNorm所需的统计量 ---
|
||||
float local_sum = 0.0f;
|
||||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||||
local_sum += x[base_offset + i];
|
||||
}
|
||||
float total_sum = blockReduceSum(local_sum, tid);
|
||||
__shared__ float mean_val;
|
||||
if (tid == 0) mean_val = total_sum / spatial_size;
|
||||
__syncthreads();
|
||||
|
||||
float local_sum_sq = 0.0f;
|
||||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||||
float val = x[base_offset + i] - mean_val;
|
||||
local_sum_sq += val * val;
|
||||
}
|
||||
float total_sum_sq = blockReduceSum(local_sum_sq, tid);
|
||||
__shared__ float inv_std_val;
|
||||
if (tid == 0) inv_std_val = rsqrtf(total_sum_sq / spatial_size + eps);
|
||||
__syncthreads();
|
||||
|
||||
// --- 融合步骤2：计算归一化后特征的Minkowski范数 ---
|
||||
float local_minkowski_sum = 0.0f;
|
||||
for (int i = tid; i < spatial_size; i += blockDim.x) {
|
||||
float normed_val = (x[base_offset + i] - mean_val) * inv_std_val;
|
||||
local_minkowski_sum += powf(fabsf(normed_val), p);
|
||||
}
|
||||
float total_minkowski_sum = blockReduceSum(local_minkowski_sum, tid);
|
||||
|
||||
// 第一个线程计算最终结果并写入输出
|
||||
if (tid == 0) {
|
||||
float minkowski_val = powf(total_minkowski_sum, inv_p);
|
||||
int output_offset = sample_idx * C + channel_idx;
|
||||
y[output_offset] = minkowski_val;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor instancenorm_minkowski_cuda(torch::Tensor x, float p, float eps) {
|
||||
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
|
||||
TORCH_CHECK(x.dim() == 4, "X must be 4D");
|
||||
TORCH_CHECK(p > 0, "p must be positive");
|
||||
|
||||
auto x_contig = x.contiguous();
|
||||
|
||||
int N = x_contig.size(0);
|
||||
int C = x_contig.size(1);
|
||||
int H = x_contig.size(2);
|
||||
int W = x_contig.size(3);
|
||||
|
||||
// 输出是 [N, C]
|
||||
auto y = torch::zeros({N, C}, x.options());
|
||||
|
||||
float inv_p = 1.0f / p;
|
||||
const int block_size = 256;
|
||||
size_t shared_mem = block_size * sizeof(float);
|
||||
|
||||
dim3 blocks(N, C);
|
||||
|
||||
instancenorm_minkowski_kernel<<<blocks, block_size, shared_mem>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
N, C, H, W,
|
||||
p, inv_p, eps
|
||||
);
|
||||
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
instancenorm_minkowski_cpp_source = """
|
||||
torch::Tensor instancenorm_minkowski_cuda(torch::Tensor x, float p, float eps);
|
||||
"""
|
||||
|
||||
instancenorm_minkowski = load_inline(
|
||||
name="instancenorm_minkowski",
|
||||
cpp_sources=instancenorm_minkowski_cpp_source,
|
||||
cuda_sources=instancenorm_minkowski_source,
|
||||
functions=["instancenorm_minkowski_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, num_features=64, p=2, eps=1e-5):
|
||||
super(ModelNew, self).__init__()
|
||||
self.p = p
|
||||
self.eps = eps
|
||||
self.instancenorm_minkowski = instancenorm_minkowski
|
||||
|
||||
def forward(self, x):
|
||||
return self.instancenorm_minkowski.instancenorm_minkowski_cuda(x, self.p, self.eps)
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Instance Normalization followed by a Minkowski Transform.
|
||||
This version uses standard PyTorch operations for a fair baseline.
|
||||
"""
|
||||
def __init__(self, num_features=64, p=2, eps=1e-5):
|
||||
super(Model, self).__init__()
|
||||
self.p = p
|
||||
if p <= 0:
|
||||
raise ValueError("p must be positive")
|
||||
|
||||
# InstanceNorm2d operates on [N, C, H, W]
|
||||
self.instance_norm = nn.InstanceNorm2d(num_features, eps=eps, affine=False, track_running_stats=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the InstanceNorm of x, then compute the Minkowski transform of the result.
|
||||
Input: [N, C, H, W]
|
||||
Output: [N, C]
|
||||
"""
|
||||
# Input validation
|
||||
if x.dim() != 4:
|
||||
raise ValueError(f"Input tensor must be 4D, got {x.dim()}D")
|
||||
|
||||
# Step 1: Apply Instance Normalization
|
||||
# This will now work correctly as H*W > 1
|
||||
normed_x = self.instance_norm(x) # Shape: [N, C, H, W]
|
||||
|
||||
# Step 2: Compute Minkowski transform per channel on the normalized features
|
||||
N, C, H, W = normed_x.shape
|
||||
normed_x_flat = normed_x.view(N, C, H * W)
|
||||
abs_normed_flat = torch.abs(normed_x_flat)
|
||||
|
||||
if self.p == 1:
|
||||
# L1 norm
|
||||
minkowski_vals = torch.sum(abs_normed_flat, dim=2) # Shape: [N, C]
|
||||
elif self.p == 2:
|
||||
# L2 norm
|
||||
minkowski_vals = torch.sqrt(torch.sum(abs_normed_flat ** 2, dim=2)) # Shape: [N, C]
|
||||
else:
|
||||
# General Lp norm
|
||||
minkowski_vals = torch.pow(torch.sum(torch.pow(abs_normed_flat, self.p), dim=2), 1.0/self.p) # Shape: [N, C]
|
||||
|
||||
return minkowski_vals
|
||||
|
||||
# 参数配置
|
||||
batch_size = 32
|
||||
num_features = 64
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, num_features, height, width)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return [num_features, 2] # num_features and p value
|
||||
|
|
@ -0,0 +1,158 @@
|
|||
You are an expert CUDA programmer tasked with accelerating a PyTorch model by replacing its operators with a highly optimized, custom CUDA kernel. You should consider operator fusion and algorithmic optimizations to achieve maximum speedup.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple ReLU:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def init(self) -> None:
|
||||
super().init()
|
||||
|
||||
def forward(self, x):
|
||||
return torch.relu(x)
|
||||
def get_inputs():
|
||||
x = torch.randn(1, 128).cuda()
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
||||
|
||||
|
||||
The example new architecture with a custom CUDA kernel looks like this:
|
||||
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
relu_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
global void relu_kernel(const float* x, float* y, int size) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < size) {
|
||||
y[idx] = fmaxf(x[idx], 0.f);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor relu_cuda(torch::Tensor x) {
|
||||
auto size = x.numel();
|
||||
auto y = torch::empty_like(x);
|
||||
const int block_size = 256;
|
||||
int num_blocks = (size + block_size - 1) / block_size;
|
||||
relu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), size);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
relu_cpp_source = """
|
||||
torch::Tensor relu_cuda(torch::Tensor x);
|
||||
"""
|
||||
|
||||
Compile the inline CUDA code
|
||||
relu = load_inline(
|
||||
name=“relu”,
|
||||
cpp_sources=relu_cpp_source,
|
||||
cuda_sources=relu_source,
|
||||
functions=[“relu_cuda”],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def init(self):
|
||||
super(ModelNew, self).init()
|
||||
self.relu = relu # The module containing the kernel
|
||||
|
||||
def forward(self, x):
|
||||
return self.relu.relu_cuda(x)
|
||||
def get_inputs():
|
||||
x = torch.randn(1, 128).cuda()
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
Now, you are given the following PyTorch architecture to accelerate. The model first applies Instance Normalization to a 4D input tensor and then computes the Minkowski norm of each channel in the normalized output. This baseline implementation uses standard PyTorch operations.
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
“”"
|
||||
Instance Normalization followed by a Minkowski Transform.
|
||||
This version uses standard PyTorch operations for a fair baseline.
|
||||
“”"
|
||||
def init(self, num_features=64, p=2, eps=1e-5):
|
||||
super(Model, self).init()
|
||||
self.p = p
|
||||
if p <= 0:
|
||||
raise ValueError(“p must be positive”)
|
||||
|
||||
# InstanceNorm2d operates on [N, C, H, W]
|
||||
self.instance_norm = nn.InstanceNorm2d(num_features, eps=eps, affine=False, track_running_stats=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the InstanceNorm of x, then compute the Minkowski transform of the result.
|
||||
Input: [N, C, H, W]
|
||||
Output: [N, C]
|
||||
"""
|
||||
# Input validation
|
||||
if x.dim() != 4:
|
||||
raise ValueError(f"Input tensor must be 4D, got {x.dim()}D")
|
||||
|
||||
# Step 1: Apply Instance Normalization
|
||||
normed_x = self.instance_norm(x) # Shape: [N, C, H, W]
|
||||
|
||||
# Step 2: Compute Minkowski transform per channel on the normalized features
|
||||
N, C, H, W = normed_x.shape
|
||||
normed_x_flat = normed_x.view(N, C, H * W)
|
||||
abs_normed_flat = torch.abs(normed_x_flat)
|
||||
|
||||
if self.p == 1:
|
||||
minkowski_vals = torch.sum(abs_normed_flat, dim=2) # Shape: [N, C]
|
||||
elif self.p == 2:
|
||||
minkowski_vals = torch.sqrt(torch.sum(abs_normed_flat ** 2, dim=2)) # Shape: [N, C]
|
||||
else:
|
||||
minkowski_vals = torch.pow(torch.sum(torch.pow(abs_normed_flat, self.p), dim=2), 1.0/self.p) # Shape: [N, C]
|
||||
|
||||
return minkowski_vals
|
||||
参数配置
|
||||
batch_size = 32
|
||||
num_features = 64
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, num_features, height, width)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
return [num_features, 2] # num_features and p value
|
||||
|
||||
Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that fuses the Instance Normalization calculation and the Minkowski norm calculation into a single kernel launch.
|
||||
|
||||
**CRITICAL REQUIREMENTS:**
|
||||
|
||||
1. **Operator Fusion:** The entire logic—computing instance mean and variance, normalizing the features, and then calculating the Minkowski norm for each channel—must be performed inside a **single CUDA kernel**. No intermediate tensors (like the normalized feature map) should be written to global memory.
|
||||
2. **Kernel Logic:**
|
||||
* Each thread block should be responsible for computing the final output for a single channel of a single sample.
|
||||
* The kernel should use a `dim3` grid where `blockIdx.x` is the sample index (N) and `blockIdx.y` is the channel index (C).
|
||||
* The kernel must perform three main stages:
|
||||
a. Calculate the `mean` and `inv_std` (1/sqrt(var + eps)) for the target channel. This requires two passes over the data (or a more complex single-pass algorithm) with efficient block-wide reductions using `extern __shared__`.
|
||||
b. Use the calculated `mean` and `inv_std` to normalize each element of the channel.
|
||||
c. Calculate the Minkowski norm of the normalized channel elements. This also requires a block-wide reduction.
|
||||
3. **Efficient Reduction:** You must implement an efficient block-wide reduction function (e.g., `blockReduceSum`) using shared memory to compute the sums required for mean, variance, and the Minkowski norm.
|
||||
4. **Final Calculation:** Inside the kernel, after the Minkowski norm is computed by the first thread (`tid == 0`), the final result should be written to the output tensor. The output tensor should have shape `[N, C]`.
|
||||
5. **Performance Optimization:** The host-side function should launch the kernel with an appropriate number of threads per block (e.g., 256) and allocate the required shared memory.
|
||||
6. **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[num_features, 2]` to match the baseline.
|
||||
7. **No Fast Math:** Do not use `--use_fast_math` in the compilation flags to ensure numerical accuracy.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from minkowski_instancenorm_torchcode import Model, get_inputs, get_init_inputs
|
||||
from minkowski_instancenorm_cudacode 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 minkowski_instancenorm 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA minkowski_instancenorm 平均执行时间: {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