forked from ccf-ai-infra/GPUCodeForces
110 lines
3.0 KiB
Python
110 lines
3.0 KiB
Python
import torch
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
logbeta_source = """
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
#include <cmath>
|
||
|
||
// VECTORIZED版本：4元素向量化处理（最终优化版本）
|
||
__global__ void logbeta_vectorized_kernel(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ y,
|
||
float* __restrict__ z,
|
||
int size
|
||
) {
|
||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||
int stride = blockDim.x * gridDim.x;
|
||
|
||
// 每次处理4个元素
|
||
int vec_size = size / 4;
|
||
|
||
for (int i = tid; i < vec_size; i += stride) {
|
||
int base_idx = i * 4;
|
||
|
||
// 加载4个元素对
|
||
float x0 = x[base_idx + 0];
|
||
float x1 = x[base_idx + 1];
|
||
float x2 = x[base_idx + 2];
|
||
float x3 = x[base_idx + 3];
|
||
|
||
float y0 = y[base_idx + 0];
|
||
float y1 = y[base_idx + 1];
|
||
float y2 = y[base_idx + 2];
|
||
float y3 = y[base_idx + 3];
|
||
|
||
// 计算4个log beta值（使用标准库函数确保精度）
|
||
z[base_idx + 0] = lgammaf(x0) + lgammaf(y0) - lgammaf(x0 + y0);
|
||
z[base_idx + 1] = lgammaf(x1) + lgammaf(y1) - lgammaf(x1 + y1);
|
||
z[base_idx + 2] = lgammaf(x2) + lgammaf(y2) - lgammaf(x2 + y2);
|
||
z[base_idx + 3] = lgammaf(x3) + lgammaf(y3) - lgammaf(x3 + y3);
|
||
}
|
||
|
||
// 处理剩余元素
|
||
int remainder = size % 4;
|
||
if (tid == 0 && remainder > 0) {
|
||
int start_idx = vec_size * 4;
|
||
for (int i = start_idx; i < size; i++) {
|
||
float xi = x[i];
|
||
float yi = y[i];
|
||
z[i] = lgammaf(xi) + lgammaf(yi) - lgammaf(xi + yi);
|
||
}
|
||
}
|
||
}
|
||
|
||
torch::Tensor logbeta_cuda(
|
||
torch::Tensor x,
|
||
torch::Tensor y,
|
||
std::string mode = "vectorized"
|
||
) {
|
||
auto x_contig = x.contiguous();
|
||
auto y_contig = y.contiguous();
|
||
auto z = torch::empty_like(x_contig);
|
||
int size = x_contig.numel();
|
||
|
||
if (mode == "vectorized") {
|
||
// 向量化版本
|
||
const int block_size = 256;
|
||
int num_blocks = (size / 4 + block_size - 1) / block_size;
|
||
int grid_size = std::min(num_blocks, 65535);
|
||
|
||
logbeta_vectorized_kernel<<<grid_size, block_size>>>(
|
||
x_contig.data_ptr<float>(),
|
||
y_contig.data_ptr<float>(),
|
||
z.data_ptr<float>(),
|
||
size
|
||
);
|
||
}
|
||
|
||
return z;
|
||
}
|
||
"""
|
||
|
||
logbeta_cpp_source = """
|
||
torch::Tensor logbeta_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
|
||
"""
|
||
|
||
# 编译CUDA代码
|
||
logbeta = load_inline(
|
||
name="logbeta_vectorized",
|
||
cpp_sources=logbeta_cpp_source,
|
||
cuda_sources=logbeta_source,
|
||
functions=["logbeta_cuda"],
|
||
extra_cuda_cflags=[
|
||
"-O3",
|
||
"--use_fast_math",
|
||
"-std=c++17",
|
||
"-maxrregcount=64"
|
||
],
|
||
verbose=True
|
||
)
|
||
|
||
class ModelNew(torch.nn.Module):
|
||
def __init__(self, mode="vectorized"):
|
||
super(ModelNew, self).__init__()
|
||
self.mode = mode
|
||
self.logbeta = logbeta # The module containing the kernel
|
||
|
||
def forward(self, x, y):
|
||
return self.logbeta.logbeta_cuda(x, y, self.mode)
|