forked from ccf-ai-infra/GPUCodeForces
finish MinkowskiDistance #7
This commit is contained in:
parent
10eed82956
commit
2c4ad87dc2
|
|
@ -0,0 +1,148 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
EPS = 1e-6
|
||||
|
||||
assert (C * H * W) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, p=3.0):
|
||||
super().__init__()
|
||||
self.p = float(p)
|
||||
self.eps = EPS
|
||||
self.blocks_per_instance = 16
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
void minkowski_sum_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor y,
|
||||
torch::Tensor out,
|
||||
float p,
|
||||
int N,
|
||||
int D,
|
||||
int blocks_per_instance);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#define WARP_SIZE 32
|
||||
|
||||
__inline__ __device__ float warp_reduce_sum(float val) {
|
||||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float block_reduce_sum(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % WARP_SIZE;
|
||||
int wid = threadIdx.x / WARP_SIZE;
|
||||
|
||||
val = warp_reduce_sum(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : 0.0f;
|
||||
if (wid == 0) val = warp_reduce_sum(val);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void minkowski_sum_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ out,
|
||||
float p,
|
||||
int N,
|
||||
int D_vec
|
||||
) {
|
||||
int n_idx = blockIdx.x;
|
||||
int stride = blockDim.x * gridDim.y;
|
||||
int d_start = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
int offset_base = n_idx * D_vec * 4;
|
||||
|
||||
const float4* x4_ptr = reinterpret_cast<const float4*>(x + offset_base);
|
||||
const float4* y4_ptr = reinterpret_cast<const float4*>(y + offset_base);
|
||||
|
||||
float local_sum = 0.0f;
|
||||
|
||||
for (int i = d_start; i < D_vec; i += stride) {
|
||||
float4 vx = x4_ptr[i];
|
||||
float4 vy = y4_ptr[i];
|
||||
|
||||
local_sum += powf(fabsf(vx.x - vy.x), p);
|
||||
local_sum += powf(fabsf(vx.y - vy.y), p);
|
||||
local_sum += powf(fabsf(vx.z - vy.z), p);
|
||||
local_sum += powf(fabsf(vx.w - vy.w), p);
|
||||
}
|
||||
|
||||
float block_sum = block_reduce_sum(local_sum);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
atomicAdd(&out[n_idx], block_sum);
|
||||
}
|
||||
}
|
||||
|
||||
void minkowski_sum_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor y,
|
||||
torch::Tensor out,
|
||||
float p,
|
||||
int N,
|
||||
int D,
|
||||
int blocks_per_instance)
|
||||
{
|
||||
dim3 blocks(N, blocks_per_instance);
|
||||
dim3 threads(256);
|
||||
int D_vec = D / 4;
|
||||
|
||||
minkowski_sum_kernel<<<blocks, threads>>>(
|
||||
x.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
out.data_ptr<float>(),
|
||||
p,
|
||||
N,
|
||||
D_vec
|
||||
);
|
||||
}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name='minkowski_cuda_opt_v2',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['minkowski_sum_cuda'],
|
||||
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
if not x.is_cuda: x = x.cuda()
|
||||
if not y.is_cuda: y = y.cuda()
|
||||
|
||||
x = x.contiguous()
|
||||
y = y.contiguous()
|
||||
|
||||
N, C, H, W = x.size()
|
||||
D = C * H * W
|
||||
|
||||
out = torch.zeros(N, device=x.device, dtype=x.dtype)
|
||||
|
||||
self.op.minkowski_sum_cuda(
|
||||
x,
|
||||
y,
|
||||
out,
|
||||
self.p,
|
||||
N,
|
||||
D,
|
||||
self.blocks_per_instance
|
||||
)
|
||||
|
||||
return torch.pow(out + self.eps, 1.0 / self.p)
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
class MinkowskiDistance(nn.Module):
|
||||
|
||||
def __init__(self, p=2.0, keepdim=False, eps=1e-6):
|
||||
super().__init__()
|
||||
self.p = p
|
||||
self.keepdim = keepdim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
diff = torch.abs(x - y)
|
||||
|
||||
pow_diff = torch.pow(diff, self.p)
|
||||
|
||||
sum_pow = torch.sum(pow_diff, dim=[1, 2, 3], keepdim=self.keepdim)
|
||||
|
||||
output = torch.pow(sum_pow + self.eps, 1.0 / self.p)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, p=3.0):
|
||||
super().__init__()
|
||||
self.op = MinkowskiDistance(p=p, keepdim=False, eps=EPS)
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return self.op(x, y)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
y = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [x, y]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
p = 3.0
|
||||
return [p]
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Technologies Used :
|
||||
|
||||
PyTorch: Deep learning framework
|
||||
|
||||
CUDA: GPU acceleration for parallel computing
|
||||
|
||||
C++/CUDA C++: High-performance kernel programming
|
||||
|
||||
Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators
|
||||
|
||||
Minkowski Distance: Generalized distance metric with parameter p
|
||||
|
||||
Vectorized Memory Access (float4): Uses 128-bit wide loads (4 floats) to improve memory bandwidth utilization
|
||||
|
||||
Two-Dimensional Grid Layout: Uses dim3(N, blocks_per_instance) for parallel processing across batches and feature dimensions
|
||||
|
||||
Warp-Level Primitives: Uses __shfl_down_sync for efficient intra-warp reduction
|
||||
|
||||
Two-Stage Parallel Reduction: Combines warp-level reduction with shared memory and block-level reduction
|
||||
|
||||
Strided Memory Access: Threads process elements with calculated stride for load distribution
|
||||
|
||||
Atomic Operations (atomicAdd): Safely accumulates results from multiple thread blocks
|
||||
|
||||
Mathematical Operations: Uses powf, fabsf for Minkowski distance calculation
|
||||
|
||||
Fast Math Operations: Uses --use_fast_math compiler flag for optimized mathematical functions
|
||||
|
||||
Configurable Parallelism: blocks_per_instance parameter controls the degree of parallelism per sample
|
||||
|
||||
Memory Coalescing: Optimized memory access patterns through contiguous tensor layout and vectorized loads
|
||||
|
||||
Numerical Stability: Adds epsilon (eps) to prevent numerical issues in power operations
|
||||
|
||||
Flexible Distance Metric: Supports arbitrary p-values for generalized Minkowski distance
|
||||
|
||||
Tensor Contiguity Enforcement: Ensures optimal memory layout in PyTorch wrapper
|
||||
|
||||
Automatic Device Placement: Ensures tensors are on CUDA device
|
||||
|
||||
Efficient Reduction Pattern: Implements hierarchical reduction from thread to warp to block level
|
||||
|
||||
|
||||
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
|
||||
import torch.nn.functional as F
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
class MinkowskiDistance(nn.Module):
|
||||
|
||||
def __init__(self, p=2.0, keepdim=False, eps=1e-6):
|
||||
super().__init__()
|
||||
self.p = p
|
||||
self.keepdim = keepdim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
diff = torch.abs(x - y)
|
||||
|
||||
pow_diff = torch.pow(diff, self.p)
|
||||
|
||||
sum_pow = torch.sum(pow_diff, dim=[1, 2, 3], keepdim=self.keepdim)
|
||||
|
||||
output = torch.pow(sum_pow + self.eps, 1.0 / self.p)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, p=3.0):
|
||||
super().__init__()
|
||||
self.op = MinkowskiDistance(p=p, keepdim=False, eps=EPS)
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return self.op(x, y)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
y = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [x, y]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
p = 3.0
|
||||
return [p]
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import torch
|
||||
import time
|
||||
from MinkowskiDistance_torch import Model, get_inputs, get_init_inputs
|
||||
from MinkowskiDistance_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用")
|
||||
return
|
||||
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 准备输入数据
|
||||
inputs = [x.cuda(device=device) for x in get_inputs()]
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
|
||||
# 初始化模型
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
# 预热GPU
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 正式测试
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 精度验证
|
||||
abs_diff = torch.abs(output_torch - output_cuda)
|
||||
max_diff = torch.max(abs_diff).item()
|
||||
mean_diff = torch.mean(abs_diff).item()
|
||||
|
||||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||||
print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = True
|
||||
else:
|
||||
print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = False
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 预热GPU
|
||||
for _ in range(10):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 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内置Swish平均执行时间: {torch_time:.6f}秒")
|
||||
print(f"自定义CUDA Swish平均执行时间: {cuda_time:.6f}秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue