forked from ccf-ai-infra/GPUCodeForces
125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
# pairwisedistance_cuda.py
|
|
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from pairwisedistance_torch import BATCH_SIZE, FEATURE_DIM, P_NORM # 导入维度常量
|
|
|
|
EPS = 1e-6
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor pdist_forward_cuda(torch::Tensor x1, torch::Tensor x2, float eps_val);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <float.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
#define WARP_SIZE 32
|
|
#define FEATURE_DIM_VAL {FEATURE_DIM}
|
|
|
|
__device__ __forceinline__ 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;
|
|
}}
|
|
|
|
__global__ void pdist_fused_kernel(
|
|
const float* __restrict__ x1,
|
|
const float* __restrict__ x2,
|
|
float* __restrict__ y_out,
|
|
int num_batches,
|
|
int feature_dim,
|
|
float eps_val
|
|
) {{
|
|
|
|
__shared__ float s_data[BLOCK_SIZE];
|
|
|
|
int batch_idx = blockIdx.x;
|
|
if (batch_idx >= num_batches) return;
|
|
|
|
const float* a_ptr = x1 + batch_idx * feature_dim;
|
|
const float* b_ptr = x2 + batch_idx * feature_dim;
|
|
|
|
float thread_diff_sq_sum = 0.0f;
|
|
|
|
for (int i = threadIdx.x; i < feature_dim; i += blockDim.x) {{
|
|
float a_val = a_ptr[i];
|
|
float b_val = b_ptr[i];
|
|
|
|
float diff = a_val - b_val;
|
|
// L2 Norm: |diff|^2
|
|
thread_diff_sq_sum += diff * diff;
|
|
}}
|
|
|
|
thread_diff_sq_sum = warp_reduce_sum(thread_diff_sq_sum);
|
|
|
|
int warp_id = threadIdx.x / WARP_SIZE;
|
|
int lane_id = threadIdx.x % WARP_SIZE;
|
|
int num_warps = blockDim.x / WARP_SIZE;
|
|
|
|
if (lane_id == 0) {{
|
|
s_data[warp_id] = thread_diff_sq_sum;
|
|
}}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x < num_warps) {
|
|
float total_diff_sq_sum = warp_reduce_sum(s_data[threadIdx.x]);
|
|
|
|
if (threadIdx.x == 0) {
|
|
float result = sqrtf(total_diff_sq_sum + eps_val);
|
|
y_out[batch_idx] = result;
|
|
}
|
|
}
|
|
}}
|
|
|
|
torch::Tensor pdist_forward_cuda(torch::Tensor x1, torch::Tensor x2, float eps_val) {
|
|
TORCH_CHECK(x1.is_cuda(), "Input must be a CUDA tensor");
|
|
x1 = x1.contiguous();
|
|
x2 = x2.contiguous();
|
|
|
|
int num_batches = x1.size(0);
|
|
int feature_dim = x1.size(1);
|
|
|
|
auto output = torch::empty({num_batches}, x1.options());
|
|
|
|
const int block_size = BLOCK_SIZE;
|
|
const int grid_size = num_batches;
|
|
|
|
size_t shared_mem_size = (block_size / WARP_SIZE) * sizeof(float);
|
|
|
|
pdist_fused_kernel<<<grid_size, block_size, shared_mem_size>>>(
|
|
x1.data_ptr<float>(),
|
|
x2.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
num_batches,
|
|
feature_dim,
|
|
eps_val
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
self.pdist_op = load_inline(
|
|
name="pdist_fused_op_safest",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["pdist_forward_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=True
|
|
)
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
|
return self.pdist_op.pdist_forward_cuda(x1, x2, EPS) |