forked from ccf-ai-infra/GPUCodeForces
106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
canberra_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
|
|
// 每个线程块处理一个样本，块内线程并行处理特征维度
|
|
__global__ void canberra_distance_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ y,
|
|
float* __restrict__ distances,
|
|
int feature_dim
|
|
) {
|
|
// 使用动态共享内存进行块内归约
|
|
extern __shared__ float sdata[];
|
|
|
|
int tid = threadIdx.x;
|
|
int sample_idx = blockIdx.x;
|
|
|
|
// 每个线程计算自己负责的部分和
|
|
float partial_sum = 0.0f;
|
|
int base_addr = sample_idx * feature_dim;
|
|
|
|
// 循环处理所有特征，步长为线程数
|
|
for (int i = tid; i < feature_dim; i += blockDim.x) {
|
|
float x_val = x[base_addr + i];
|
|
float y_val = y[base_addr + i];
|
|
|
|
float abs_diff = fabsf(x_val - y_val);
|
|
float abs_sum = fabsf(x_val) + fabsf(y_val);
|
|
|
|
if (abs_sum > 0.0f) {
|
|
partial_sum += abs_diff / abs_sum;
|
|
}
|
|
}
|
|
|
|
sdata[tid] = partial_sum;
|
|
__syncthreads();
|
|
|
|
// 标准的并行归约求和
|
|
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
|
if (tid < stride) {
|
|
sdata[tid] += sdata[tid + stride];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
// 线程0将最终结果写回
|
|
if (tid == 0) {
|
|
distances[sample_idx] = sdata[0];
|
|
}
|
|
}
|
|
|
|
torch::Tensor canberra_cuda(torch::Tensor x, torch::Tensor y) {
|
|
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
|
|
TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
|
|
TORCH_CHECK(x.sizes() == y.sizes(), "X and Y must have the same shape");
|
|
|
|
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 distances = torch::zeros({batch_size}, x.options());
|
|
|
|
const int block_size = BLOCK_SIZE;
|
|
size_t shared_mem = block_size * sizeof(float);
|
|
|
|
canberra_distance_kernel<<<batch_size, block_size, shared_mem>>>(
|
|
x_contig.data_ptr<float>(),
|
|
y_contig.data_ptr<float>(),
|
|
distances.data_ptr<float>(),
|
|
feature_dim
|
|
);
|
|
|
|
return distances;
|
|
}
|
|
"""
|
|
|
|
canberra_cpp_source = """
|
|
torch::Tensor canberra_cuda(torch::Tensor x, torch::Tensor y);
|
|
"""
|
|
|
|
# 编译CUDA代码
|
|
canberra = load_inline(
|
|
name="canberra_fixed",
|
|
cpp_sources=canberra_cpp_source,
|
|
cuda_sources=canberra_source,
|
|
functions=["canberra_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=True
|
|
)
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self):
|
|
super(ModelNew, self).__init__()
|
|
self.canberra = canberra
|
|
|
|
def forward(self, x, y):
|
|
return self.canberra.canberra_cuda(x, y)
|