forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish JaccardSimilarity #9' (#128) from uucoco/GPUCodeForces:uucoco9 into main
This commit is contained in:
commit
53c068614f
|
|
@ -0,0 +1,274 @@
|
|||
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):
|
||||
super().__init__()
|
||||
self.block_size = 512
|
||||
self.eps = EPS
|
||||
|
||||
self.register_buffer('temp_buffer', torch.zeros((N, 3), dtype=torch.float32))
|
||||
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
void jaccard_fused_cuda(
|
||||
torch::Tensor x1,
|
||||
torch::Tensor x2,
|
||||
torch::Tensor temp_buffer, // [N, 3]
|
||||
int N,
|
||||
int D);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
#define BLOCK_SIZE {self.block_size}
|
||||
#define WARP_SIZE 32
|
||||
#define ILP 4
|
||||
|
||||
|
||||
__inline__ __device__ float warp_reduce_sum(float val) {{
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {{
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}}
|
||||
return val;
|
||||
}}
|
||||
|
||||
|
||||
__global__ __launch_bounds__(BLOCK_SIZE)
|
||||
void jaccard_split_kernel(
|
||||
const float* __restrict__ x1,
|
||||
const float* __restrict__ x2,
|
||||
float* __restrict__ temp_buffer, // [N, 3]
|
||||
int D_vec_total // D / 4
|
||||
) {{
|
||||
|
||||
const int n_idx = blockIdx.y; // Batch Index
|
||||
const int split_idx = blockIdx.x; // Split Index
|
||||
const int num_splits = gridDim.x;
|
||||
|
||||
|
||||
const int chunk_size = (D_vec_total + num_splits - 1) / num_splits;
|
||||
const int start_idx = split_idx * chunk_size;
|
||||
const int end_idx = min(start_idx + chunk_size, D_vec_total);
|
||||
|
||||
if (start_idx >= D_vec_total) return;
|
||||
|
||||
|
||||
const int64_t batch_offset = (int64_t)n_idx * D_vec_total * 4;
|
||||
|
||||
const float4* curr_x1 = reinterpret_cast<const float4*>(x1 + batch_offset) + start_idx + threadIdx.x;
|
||||
const float4* curr_x2 = reinterpret_cast<const float4*>(x2 + batch_offset) + start_idx + threadIdx.x;
|
||||
const float4* end_ptr = reinterpret_cast<const float4*>(x1 + batch_offset) + end_idx;
|
||||
|
||||
|
||||
float acc_dot[ILP];
|
||||
float acc_sq1[ILP];
|
||||
float acc_sq2[ILP];
|
||||
|
||||
#pragma unroll
|
||||
for (int k=0; k<ILP; ++k) {{
|
||||
acc_dot[k] = 0.0f;
|
||||
acc_sq1[k] = 0.0f;
|
||||
acc_sq2[k] = 0.0f;
|
||||
}}
|
||||
|
||||
const int stride = BLOCK_SIZE * ILP;
|
||||
|
||||
|
||||
while (curr_x1 + (ILP - 1) * BLOCK_SIZE < end_ptr) {{
|
||||
float4 r1[ILP];
|
||||
float4 r2[ILP];
|
||||
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < ILP; ++k) {{
|
||||
r1[k] = __ldg(curr_x1 + k * BLOCK_SIZE);
|
||||
r2[k] = __ldg(curr_x2 + k * BLOCK_SIZE);
|
||||
}}
|
||||
|
||||
// Compute (FMA 优化)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < ILP; ++k) {{
|
||||
// Dot: x1 * x2
|
||||
acc_dot[k] += r1[k].x * r2[k].x;
|
||||
acc_dot[k] += r1[k].y * r2[k].y;
|
||||
acc_dot[k] += r1[k].z * r2[k].z;
|
||||
acc_dot[k] += r1[k].w * r2[k].w;
|
||||
|
||||
// Sq1: x1 * x1
|
||||
acc_sq1[k] += r1[k].x * r1[k].x;
|
||||
acc_sq1[k] += r1[k].y * r1[k].y;
|
||||
acc_sq1[k] += r1[k].z * r1[k].z;
|
||||
acc_sq1[k] += r1[k].w * r1[k].w;
|
||||
|
||||
// Sq2: x2 * x2
|
||||
acc_sq2[k] += r2[k].x * r2[k].x;
|
||||
acc_sq2[k] += r2[k].y * r2[k].y;
|
||||
acc_sq2[k] += r2[k].z * r2[k].z;
|
||||
acc_sq2[k] += r2[k].w * r2[k].w;
|
||||
}}
|
||||
|
||||
curr_x1 += stride;
|
||||
curr_x2 += stride;
|
||||
}}
|
||||
|
||||
|
||||
while (curr_x1 < end_ptr) {{
|
||||
float4 v1 = __ldg(curr_x1);
|
||||
float4 v2 = __ldg(curr_x2);
|
||||
|
||||
acc_dot[0] += v1.x * v2.x + v1.y * v2.y + v1.z * v2.z + v1.w * v2.w;
|
||||
acc_sq1[0] += v1.x * v1.x + v1.y * v1.y + v1.z * v1.z + v1.w * v1.w;
|
||||
acc_sq2[0] += v2.x * v2.x + v2.y * v2.y + v2.z * v2.z + v2.w * v2.w;
|
||||
|
||||
curr_x1 += BLOCK_SIZE;
|
||||
curr_x2 += BLOCK_SIZE;
|
||||
}}
|
||||
|
||||
|
||||
float sum_dot = 0.0f;
|
||||
float sum_sq1 = 0.0f;
|
||||
float sum_sq2 = 0.0f;
|
||||
|
||||
#pragma unroll
|
||||
for (int k=0; k<ILP; ++k) {{
|
||||
sum_dot += acc_dot[k];
|
||||
sum_sq1 += acc_sq1[k];
|
||||
sum_sq2 += acc_sq2[k];
|
||||
}}
|
||||
|
||||
// 8. Warp 归约
|
||||
sum_dot = warp_reduce_sum(sum_dot);
|
||||
sum_sq1 = warp_reduce_sum(sum_sq1);
|
||||
sum_sq2 = warp_reduce_sum(sum_sq2);
|
||||
|
||||
// 9. Block 归约 (使用 Shared Memory)
|
||||
__shared__ float s_dot[BLOCK_SIZE / WARP_SIZE];
|
||||
__shared__ float s_sq1[BLOCK_SIZE / WARP_SIZE];
|
||||
__shared__ float s_sq2[BLOCK_SIZE / WARP_SIZE];
|
||||
|
||||
const int lane_id = threadIdx.x % WARP_SIZE;
|
||||
const int warp_id = threadIdx.x / WARP_SIZE;
|
||||
|
||||
if (lane_id == 0) {{
|
||||
s_dot[warp_id] = sum_dot;
|
||||
s_sq1[warp_id] = sum_sq1;
|
||||
s_sq2[warp_id] = sum_sq2;
|
||||
}}
|
||||
__syncthreads();
|
||||
|
||||
if (warp_id == 0) {{
|
||||
float block_dot = 0.0f;
|
||||
float block_sq1 = 0.0f;
|
||||
float block_sq2 = 0.0f;
|
||||
|
||||
if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
|
||||
block_dot = s_dot[lane_id];
|
||||
block_sq1 = s_sq1[lane_id];
|
||||
block_sq2 = s_sq2[lane_id];
|
||||
}}
|
||||
|
||||
block_dot = warp_reduce_sum(block_dot);
|
||||
block_sq1 = warp_reduce_sum(block_sq1);
|
||||
block_sq2 = warp_reduce_sum(block_sq2);
|
||||
|
||||
// 10. 原子累加到全局 Temp Buffer
|
||||
// Buffer Layout: [N, 3] -> Row n: [dot, sq1, sq2]
|
||||
if (lane_id == 0) {{
|
||||
float* dst = temp_buffer + n_idx * 3;
|
||||
atomicAdd(&dst[0], block_dot);
|
||||
atomicAdd(&dst[1], block_sq1);
|
||||
atomicAdd(&dst[2], block_sq2);
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
|
||||
void jaccard_fused_cuda(
|
||||
torch::Tensor x1,
|
||||
torch::Tensor x2,
|
||||
torch::Tensor temp_buffer,
|
||||
int N,
|
||||
int D)
|
||||
{{
|
||||
int D_vec = D / 4;
|
||||
|
||||
int device_id;
|
||||
cudaGetDevice(&device_id);
|
||||
int sm_count;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id);
|
||||
|
||||
|
||||
int target_blocks = sm_count * 4;
|
||||
int splits = (target_blocks + N - 1) / N;
|
||||
|
||||
|
||||
int max_splits = (D_vec + 512 - 1) / 512;
|
||||
if (splits > max_splits) splits = max_splits;
|
||||
if (splits < 1) splits = 1;
|
||||
|
||||
if (splits > 512) splits = 512;
|
||||
|
||||
dim3 blocks(splits, N);
|
||||
dim3 threads(BLOCK_SIZE);
|
||||
|
||||
jaccard_split_kernel<<<blocks, threads>>>(
|
||||
x1.data_ptr<float>(),
|
||||
x2.data_ptr<float>(),
|
||||
temp_buffer.data_ptr<float>(),
|
||||
D_vec
|
||||
);
|
||||
}}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name='jaccard_cuda_opt_v1',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['jaccard_fused_cuda'],
|
||||
extra_cuda_cflags=[
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'-Xptxas=-v'
|
||||
],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
if not x1.is_contiguous(): x1 = x1.contiguous()
|
||||
if not x2.is_contiguous(): x2 = x2.contiguous()
|
||||
|
||||
N, C, H, W = x1.size()
|
||||
D = C * H * W
|
||||
|
||||
self.temp_buffer.zero_()
|
||||
|
||||
self.op.jaccard_fused_cuda(
|
||||
x1,
|
||||
x2,
|
||||
self.temp_buffer,
|
||||
N,
|
||||
D
|
||||
)
|
||||
|
||||
dot = self.temp_buffer[:, 0]
|
||||
sq1 = self.temp_buffer[:, 1]
|
||||
sq2 = self.temp_buffer[:, 2]
|
||||
|
||||
union = sq1 + sq2 - dot
|
||||
return dot / (union + self.eps)
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
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 JaccardSimilarity(nn.Module):
|
||||
|
||||
def __init__(self, eps=1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
intersection = torch.sum(x1 * x2, dim=[1, 2, 3])
|
||||
|
||||
sum_sq1 = torch.sum(x1 * x1, dim=[1, 2, 3])
|
||||
sum_sq2 = torch.sum(x2 * x2, dim=[1, 2, 3])
|
||||
|
||||
union = sum_sq1 + sum_sq2 - intersection
|
||||
|
||||
return intersection / (union + self.eps)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.op = JaccardSimilarity(EPS)
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
return self.op(x1, x2)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x1 = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
x2 = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [x1, x2]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,101 @@
|
|||
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
|
||||
|
||||
Jaccard Similarity (Intersection over Union): Similarity measure between sets or vectors
|
||||
|
||||
Vectorized Memory Access (float4): Uses 128-bit wide loads (4 floats) to improve memory bandwidth utilization
|
||||
|
||||
Instruction-Level Parallelism (ILP): Processes multiple float4 elements per loop iteration to hide instruction latency
|
||||
|
||||
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
|
||||
|
||||
Multi-Dimensional Grid Layout: Uses dim3(splits, N) for parallel processing across splits and batches
|
||||
|
||||
Dynamic Kernel Configuration: Calculates optimal split count based on GPU SM count and data size
|
||||
|
||||
Fused Kernel Design: Computes dot product and squared sums simultaneously in single kernel
|
||||
|
||||
Atomic Operations (atomicAdd): Safely accumulates results from multiple thread blocks to temporary buffer
|
||||
|
||||
Temporary Buffer Strategy: Uses pre-allocated buffer [N, 3] to store intermediate results (dot, sq1, sq2)
|
||||
|
||||
Constant Memory/__ldg: Uses read-only data cache for improved memory access patterns
|
||||
|
||||
Fast Math Operations: Uses FMA operations with --use_fast_math compiler flag
|
||||
|
||||
Memory Coalescing: Optimized memory access patterns through contiguous tensor layout
|
||||
|
||||
Pointer Chasing Loop: Efficient main loop with ILP-unrolled memory access patterns
|
||||
|
||||
Tail Processing: Handles remaining elements after main vectorized loop
|
||||
|
||||
Numerical Stability: Adds epsilon (eps) to prevent division by zero in final calculation
|
||||
|
||||
Buffer Zeroing: Clears temporary buffer before each forward pass
|
||||
|
||||
Device Query API: Uses cudaGetDevice and cudaDeviceGetAttribute for optimal kernel configuration
|
||||
|
||||
Three-Accumulator Design: Maintains separate accumulators for dot product, x1 squared, and x2 squared
|
||||
|
||||
Efficient Union Calculation: Computes Jaccard similarity using algebraic identity: union = sum(sq1) + sum(sq2) - sum(dot)
|
||||
|
||||
Shared Memory for Warp Results: Uses separate shared memory arrays for each reduction variable
|
||||
|
||||
Boundary Checking: Handles data size variations and split boundaries safely
|
||||
|
||||
|
||||
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 JaccardSimilarity(nn.Module):
|
||||
|
||||
def __init__(self, eps=1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
intersection = torch.sum(x1 * x2, dim=[1, 2, 3])
|
||||
|
||||
sum_sq1 = torch.sum(x1 * x1, dim=[1, 2, 3])
|
||||
sum_sq2 = torch.sum(x2 * x2, dim=[1, 2, 3])
|
||||
|
||||
union = sum_sq1 + sum_sq2 - intersection
|
||||
|
||||
return intersection / (union + self.eps)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.op = JaccardSimilarity(EPS)
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
return self.op(x1, x2)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x1 = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
x2 = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [x1, x2]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from JaccardSimilarity_torch import Model, get_inputs, get_init_inputs
|
||||
from JaccardSimilarity_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