forked from ccf-ai-infra/GPUCodeForces
feat:add minkowski+leakyrelu #75
This commit is contained in:
parent
f989885dde
commit
7fcd64b831
|
|
@ -0,0 +1,174 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
minkowski_leakyrelu_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
// 通用融合内核 - Minkowski Distance + LeakyReLU
|
||||
__global__ void minkowski_leakyrelu_kernel_general(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ output,
|
||||
int batch_size,
|
||||
int feature_dim,
|
||||
float p,
|
||||
float inv_p,
|
||||
float negative_slope
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int base = sample_idx * feature_dim;
|
||||
|
||||
extern __shared__ float shared_sum[];
|
||||
shared_sum[tid] = 0.0f;
|
||||
|
||||
int stride = blockDim.x;
|
||||
for (int dim = tid; dim < feature_dim; dim += stride) {
|
||||
float diff = x[base + dim] - y[base + dim];
|
||||
float abs_diff = fabsf(diff);
|
||||
shared_sum[tid] += powf(abs_diff, p);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
shared_sum[tid] += shared_sum[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
float distance = powf(shared_sum[0], inv_p);
|
||||
// --- 融合步骤：直接计算 LeakyReLU(distance) ---
|
||||
// LeakyReLU(x) = max(x, negative_slope * x)
|
||||
output[sample_idx] = fmaxf(distance, negative_slope * distance);
|
||||
}
|
||||
}
|
||||
|
||||
// 特殊优化融合版本 - p=2 (欧几里得距离)
|
||||
__global__ void minkowski_leakyrelu_kernel_p2(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ output,
|
||||
int batch_size,
|
||||
int feature_dim,
|
||||
float negative_slope
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int base = sample_idx * feature_dim;
|
||||
float sum_sq = 0.0f;
|
||||
|
||||
for (int dim = 0; dim < feature_dim; dim++) {
|
||||
float diff = x[base + dim] - y[base + dim];
|
||||
sum_sq += diff * diff;
|
||||
}
|
||||
|
||||
float distance = sqrtf(sum_sq);
|
||||
// --- 融合步骤：直接计算 LeakyReLU(distance) ---
|
||||
output[sample_idx] = fmaxf(distance, negative_slope * distance);
|
||||
}
|
||||
|
||||
// 特殊优化融合版本 - p=1 (曼哈顿距离)
|
||||
__global__ void minkowski_leakyrelu_kernel_p1(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ output,
|
||||
int batch_size,
|
||||
int feature_dim,
|
||||
float negative_slope
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int base = sample_idx * feature_dim;
|
||||
float sum_abs = 0.0f;
|
||||
|
||||
for (int dim = 0; dim < feature_dim; dim++) {
|
||||
float diff = x[base + dim] - y[base + dim];
|
||||
sum_abs += fabsf(diff);
|
||||
}
|
||||
|
||||
// --- 融合步骤：直接计算 LeakyReLU(distance) ---
|
||||
output[sample_idx] = fmaxf(sum_abs, negative_slope * sum_abs);
|
||||
}
|
||||
|
||||
torch::Tensor minkowski_leakyrelu_cuda(torch::Tensor x, torch::Tensor y, float p, float negative_slope) {
|
||||
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
|
||||
TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
|
||||
TORCH_CHECK(p > 0, "p must be positive");
|
||||
|
||||
auto x_contig = x.contiguous();
|
||||
auto y_contig = y.contiguous();
|
||||
|
||||
int batch_size = x_contig.size(0);
|
||||
int feature_dim = x_contig.size(1);
|
||||
|
||||
auto output = torch::zeros({batch_size}, x.options());
|
||||
|
||||
if (p == 1.0f) {
|
||||
minkowski_leakyrelu_kernel_p1<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim,
|
||||
negative_slope
|
||||
);
|
||||
} else if (p == 2.0f) {
|
||||
minkowski_leakyrelu_kernel_p2<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim,
|
||||
negative_slope
|
||||
);
|
||||
} else {
|
||||
float inv_p = 1.0f / p;
|
||||
const int block_size = 256;
|
||||
size_t shared_mem = block_size * sizeof(float);
|
||||
|
||||
minkowski_leakyrelu_kernel_general<<<batch_size, block_size, shared_mem>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim,
|
||||
p,
|
||||
inv_p,
|
||||
negative_slope
|
||||
);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
minkowski_leakyrelu_cpp_source = """
|
||||
torch::Tensor minkowski_leakyrelu_cuda(torch::Tensor x, torch::Tensor y, float p, float negative_slope);
|
||||
"""
|
||||
|
||||
minkowski_leakyrelu = load_inline(
|
||||
name="minkowski_leakyrelu",
|
||||
cpp_sources=minkowski_leakyrelu_cpp_source,
|
||||
cuda_sources=minkowski_leakyrelu_source,
|
||||
functions=["minkowski_leakyrelu_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, p=2, negative_slope=0.01):
|
||||
super(ModelNew, self).__init__()
|
||||
self.p = p
|
||||
self.negative_slope = negative_slope
|
||||
self.minkowski_leakyrelu = minkowski_leakyrelu
|
||||
|
||||
def forward(self, x, y):
|
||||
return self.minkowski_leakyrelu.minkowski_leakyrelu_cuda(x, y, self.p, self.negative_slope)
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Minkowski Distance followed by a LeakyReLU activation.
|
||||
This version uses standard PyTorch operations for a fair baseline.
|
||||
"""
|
||||
def __init__(self, p=2, negative_slope=0.01):
|
||||
super(Model, self).__init__()
|
||||
self.p = p
|
||||
self.negative_slope = negative_slope
|
||||
if p <= 0:
|
||||
raise ValueError("p must be positive")
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the LeakyReLU of the Minkowski distance between x and y.
|
||||
"""
|
||||
# Input validation
|
||||
if x.shape != y.shape:
|
||||
raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
|
||||
if x.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
|
||||
|
||||
# Step 1: Compute Minkowski distance using element-wise operations
|
||||
abs_diff = torch.abs(x - y)
|
||||
|
||||
if self.p == 1:
|
||||
d = torch.sum(abs_diff, dim=1)
|
||||
elif self.p == 2:
|
||||
d = torch.sqrt(torch.sum(abs_diff ** 2, dim=1))
|
||||
else:
|
||||
d = torch.pow(torch.sum(torch.pow(abs_diff, self.p), dim=1), 1.0/self.p)
|
||||
|
||||
# Step 2: Apply LeakyReLU activation
|
||||
output = F.leaky_relu(d, negative_slope=self.negative_slope)
|
||||
|
||||
return output
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 512
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
y = torch.randn(batch_size, feature_dim)
|
||||
return [x, y]
|
||||
|
||||
def get_init_inputs():
|
||||
return [2, 0.01] # p value and negative_slope
|
||||
|
|
@ -0,0 +1,152 @@
|
|||
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 computes the Minkowski distance between two tensors and then applies a LeakyReLU activation to the resulting distances. This baseline implementation uses standard element-wise operations to ensure a clear, sample-by-sample calculation.
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
“”"
|
||||
Minkowski Distance followed by a LeakyReLU activation.
|
||||
This version uses standard PyTorch operations for a fair baseline.
|
||||
“”"
|
||||
def init(self, p=2, negative_slope=0.01):
|
||||
super(Model, self).init()
|
||||
self.p = p
|
||||
self.negative_slope = negative_slope
|
||||
if p <= 0:
|
||||
raise ValueError(“p must be positive”)
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the LeakyReLU of the Minkowski distance between x and y.
|
||||
"""
|
||||
# Input validation
|
||||
if x.shape != y.shape:
|
||||
raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
|
||||
if x.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
|
||||
|
||||
# Step 1: Compute Minkowski distance using element-wise operations
|
||||
abs_diff = torch.abs(x - y)
|
||||
|
||||
if self.p == 1:
|
||||
d = torch.sum(abs_diff, dim=1)
|
||||
elif self.p == 2:
|
||||
d = torch.sqrt(torch.sum(abs_diff ** 2, dim=1))
|
||||
else:
|
||||
d = torch.pow(torch.sum(torch.pow(abs_diff, self.p), dim=1), 1.0/self.p)
|
||||
|
||||
# Step 2: Apply LeakyReLU activation
|
||||
output = F.leaky_relu(d, negative_slope=self.negative_slope)
|
||||
|
||||
return output
|
||||
batch_size = 256
|
||||
feature_dim = 512
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
y = torch.randn(batch_size, feature_dim)
|
||||
return [x, y]
|
||||
|
||||
def get_init_inputs():
|
||||
return [2, 0.01] # p value and negative_slope
|
||||
|
||||
|
||||
|
||||
Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that fuses the Minkowski distance calculation and the LeakyReLU activation into a single kernel launch, thereby eliminating the intermediate distance tensor.
|
||||
|
||||
**CRITICAL REQUIREMENTS:**
|
||||
|
||||
1. **Operator Fusion:** The entire logic—computing the Minkowski distance for each sample in the batch and then applying the LeakyReLU activation (`max(distance, negative_slope * distance)`)—must be performed inside a **single CUDA kernel**. No intermediate distance tensors should be written to global memory.
|
||||
2. **Algorithmic Specialization:** The implementation must provide specialized, highly optimized kernels for the most common cases, `p=1` (Manhattan) and `p=2` (Euclidean), in addition to a general kernel for any `p`.
|
||||
3. **Kernel Logic:**
|
||||
* Each thread block should be responsible for computing the final output for a single sample in the batch.
|
||||
* For `p=1` and `p=2`, the kernel should be a simple loop without shared memory reduction for maximum efficiency.
|
||||
* For the general `p` case, the kernel should use a multi-threaded reduction within a thread block with `extern __shared__`.
|
||||
4. **Final Calculation:** Inside the kernel, after the distance is computed by the first thread (`tid == 0`), the LeakyReLU activation must be applied immediately using `fmaxf(distance, negative_slope * distance)` before writing the final result to the output tensor.
|
||||
5. **Performance Optimization:** The host-side function should dispatch to the appropriate specialized kernel based on the value of `p` at runtime.
|
||||
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 `[2, 0.01]` 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_leakyrelu_torchcode import Model, get_inputs, get_init_inputs
|
||||
from minkowski_leakyrelu_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_leakyrelu 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA minkowski_leakyrelu 平均执行时间: {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