forked from ccf-ai-infra/GPUCodeForces
154 lines
4.4 KiB
Python
154 lines
4.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
import math
|
|
|
|
N_BATCH = 100
|
|
D_VECTOR = 128
|
|
|
|
DIM = 1
|
|
BLOCK_SIZE = 256
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, dim=1):
|
|
super().__init__()
|
|
|
|
if dim != 1:
|
|
raise NotImplementedError("CUDA Kernel only supports dim=1 for (N,D) tensors")
|
|
self.dim = dim
|
|
self.block_size = BLOCK_SIZE
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor hamming_distance_forward_cuda(
|
|
torch::Tensor x1,
|
|
torch::Tensor x2,
|
|
int dim
|
|
);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
|
|
#define BLOCK_SIZE {self.block_size}
|
|
|
|
/*
|
|
* Hamming Distance 融合核函数
|
|
* (输入: int64_t, 输出: float)
|
|
*/
|
|
__global__ void hamming_distance_fused_kernel(
|
|
const int64_t* __restrict__ x1_data,
|
|
const int64_t* __restrict__ x2_data,
|
|
float* __restrict__ output_data,
|
|
int N,
|
|
int D
|
|
) {{
|
|
|
|
__shared__ int s_data[BLOCK_SIZE];
|
|
|
|
const int n_idx = blockIdx.x; // [0, N-1]
|
|
const int tid = threadIdx.x; // [0, BLOCK_SIZE-1]
|
|
|
|
const int64_t* p_in1 = x1_data + (int64_t)n_idx * D;
|
|
const int64_t* p_in2 = x2_data + (int64_t)n_idx * D;
|
|
|
|
|
|
int thread_sum = 0;
|
|
|
|
for (int d = tid; d < D; d += BLOCK_SIZE) {{
|
|
if (p_in1[d] != p_in2[d]) {{
|
|
thread_sum++;
|
|
}}
|
|
}}
|
|
s_data[tid] = thread_sum;
|
|
__syncthreads();
|
|
|
|
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
|
|
if (tid < offset) {{
|
|
s_data[tid] += s_data[tid + offset];
|
|
}}
|
|
__syncthreads();
|
|
}}
|
|
|
|
|
|
if (tid == 0) {{
|
|
output_data[n_idx] = (float)s_data[0];
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor hamming_distance_forward_cuda(
|
|
torch::Tensor x1,
|
|
torch::Tensor x2,
|
|
int dim
|
|
) {{
|
|
|
|
TORCH_CHECK(x1.dim() == 2, "CUDA Kernel only supports 2D (N, D) input");
|
|
TORCH_CHECK(dim == 1, "CUDA Kernel only supports dim=1");
|
|
TORCH_CHECK(x1.is_cuda() && x2.is_cuda(), "Inputs must be CUDA");
|
|
TORCH_CHECK(x1.is_contiguous() && x2.is_contiguous(), "Inputs must be contiguous");
|
|
TORCH_CHECK(x1.sizes() == x2.sizes(), "Input shapes must match");
|
|
|
|
|
|
TORCH_CHECK(x1.scalar_type() == torch::kLong, "Input 1 must be torch.long");
|
|
TORCH_CHECK(x2.scalar_type() == torch::kLong, "Input 2 must be torch.long");
|
|
|
|
const int64_t N = x1.size(0);
|
|
const int64_t D = x1.size(1);
|
|
|
|
|
|
auto output_options = torch::TensorOptions()
|
|
.dtype(torch::kFloat32)
|
|
.device(x1.device());
|
|
auto output = torch::empty({{N}}, output_options);
|
|
|
|
dim3 grid_dim(N);
|
|
dim3 block_dim(BLOCK_SIZE);
|
|
|
|
hamming_distance_fused_kernel<<<grid_dim, block_dim>>>(
|
|
x1.data_ptr<int64_t>(),
|
|
x2.data_ptr<int64_t>(),
|
|
output.data_ptr<float>(),
|
|
N, D
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = ['-O3']
|
|
|
|
# JIT (Just-In-Time) 编译
|
|
self.hamming_op = load_inline(
|
|
name="hamming_dist_op_v1",
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["hamming_distance_forward_cuda"],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
|
# 确保连续性
|
|
x1_cont = x1.contiguous()
|
|
x2_cont = x2.contiguous()
|
|
|
|
# CUDA 核函数要求 Long 类型
|
|
# (如果输入不是 Long, .to() 会创建一个副本)
|
|
x1_long = x1_cont.to(torch.long)
|
|
x2_long = x2_cont.to(torch.long)
|
|
|
|
return self.hamming_op.hamming_distance_forward_cuda(
|
|
x1_long,
|
|
x2_long,
|
|
self.dim
|
|
) |