forked from ccf-ai-infra/GPUCodeForces
216 lines
7.1 KiB
Python
216 lines
7.1 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.cpp_extension import load_inline
|
||
import os
|
||
|
||
# 定义维度常量
|
||
N, C, H, W = 32, 64, 56, 56
|
||
EPS = 1e-6
|
||
INSTANCE_SIZE = C * H * W # 200704
|
||
|
||
# 确保 C*H*W 是 4 的倍数以便向量化
|
||
assert (INSTANCE_SIZE) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"
|
||
|
||
|
||
class ModelNew(nn.Module):
|
||
"""
|
||
Tanimoto 系数的 CUDA 优化实现。
|
||
|
||
优化点:
|
||
1. 融合: 将 3 次归约 (dot, x_sq, y_sq) 和 1 次除法融合到一个 CUDA kernel。
|
||
2. 向量化: 使用 float4 进行 128 位内存访问。
|
||
3. 共享内存: 用于高效的块内归约。
|
||
4. 精度匹配: 使用 'volatile' 禁用 FMA，以 100% 匹配 PyTorch 的分步舍入。
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.eps = EPS
|
||
self._compile_cuda_kernel()
|
||
|
||
def _compile_cuda_kernel(self):
|
||
# 1. C++ 接口定义
|
||
cpp_source = """
|
||
#include <torch/extension.h>
|
||
|
||
torch::Tensor tanimoto_forward_cuda(
|
||
torch::Tensor x,
|
||
torch::Tensor y,
|
||
float eps,
|
||
int N,
|
||
int D);
|
||
"""
|
||
|
||
# 2. CUDA 内核实现
|
||
cuda_source = f"""
|
||
#include <cuda_runtime.h>
|
||
#include <device_launch_parameters.h>
|
||
#include <cmath>
|
||
|
||
// ============================================================
|
||
// Tanimoto 融合内核
|
||
// ============================================================
|
||
__global__ void tanimoto_fused_kernel(
|
||
const float* __restrict__ x,
|
||
const float* __restrict__ y,
|
||
float* __restrict__ output,
|
||
float eps,
|
||
int N, // 批次大小
|
||
int D // 实例大小 (C*H*W)
|
||
) {{
|
||
// 每个块处理一个实例 (n_idx)
|
||
const int n_idx = blockIdx.x;
|
||
if (n_idx >= N) return;
|
||
|
||
const int t_idx = threadIdx.x;
|
||
const int block_size = blockDim.x; // e.g., 256
|
||
|
||
// 指向此实例数据的开头
|
||
const int instance_offset = n_idx * D;
|
||
const float4* x4_ptr = reinterpret_cast<const float4*>(x + instance_offset);
|
||
const float4* y4_ptr = reinterpret_cast<const float4*>(y + instance_offset);
|
||
|
||
const int D_vec = D / 4; // 向量化维度
|
||
|
||
// 线程本地累加器
|
||
float thread_dot = 0.0f;
|
||
float thread_x_sq = 0.0f;
|
||
float thread_y_sq = 0.0f;
|
||
|
||
// --- 1. 线程本地向量化累加 ---
|
||
for (int i = t_idx; i < D_vec; i += block_size) {{
|
||
float4 x_val = x4_ptr[i];
|
||
float4 y_val = y4_ptr[i];
|
||
|
||
// --- 精度修复: 禁用 FMA ---
|
||
// (x * y)
|
||
volatile float m_dot_x = x_val.x * y_val.x;
|
||
volatile float m_dot_y = x_val.y * y_val.y;
|
||
volatile float m_dot_z = x_val.z * y_val.z;
|
||
volatile float m_dot_w = x_val.w * y_val.w;
|
||
|
||
// (x * x)
|
||
volatile float m_x_sq_x = x_val.x * x_val.x;
|
||
volatile float m_x_sq_y = x_val.y * x_val.y;
|
||
volatile float m_x_sq_z = x_val.z * x_val.z;
|
||
volatile float m_x_sq_w = x_val.w * x_val.w;
|
||
|
||
// (y * y)
|
||
volatile float m_y_sq_x = y_val.x * y_val.x;
|
||
volatile float m_y_sq_y = y_val.y * y_val.y;
|
||
volatile float m_y_sq_z = y_val.z * y_val.z;
|
||
volatile float m_y_sq_w = y_val.w * y_val.w;
|
||
|
||
thread_dot += m_dot_x + m_dot_y + m_dot_z + m_dot_w;
|
||
thread_x_sq += m_x_sq_x + m_x_sq_y + m_x_sq_z + m_x_sq_w;
|
||
thread_y_sq += m_y_sq_x + m_y_sq_y + m_y_sq_z + m_y_sq_w;
|
||
}}
|
||
|
||
// --- 2. 块内共享内存归约 ---
|
||
// (BLOCK_SIZE 必须是 2 的幂, e.g., 256)
|
||
extern __shared__ float s_data[];
|
||
float* s_dot = s_data;
|
||
float* s_x_sq = &s_data[block_size];
|
||
float* s_y_sq = &s_data[block_size * 2];
|
||
|
||
s_dot[t_idx] = thread_dot;
|
||
s_x_sq[t_idx] = thread_x_sq;
|
||
s_y_sq[t_idx] = thread_y_sq;
|
||
|
||
__syncthreads();
|
||
|
||
// (使用标准树形归约)
|
||
for (int s = block_size / 2; s > 0; s >>= 1) {{
|
||
if (t_idx < s) {{
|
||
s_dot[t_idx] += s_dot[t_idx + s];
|
||
s_x_sq[t_idx] += s_x_sq[t_idx + s];
|
||
s_y_sq[t_idx] += s_y_sq[t_idx + s];
|
||
}}
|
||
__syncthreads();
|
||
}}
|
||
|
||
// --- 3. 线程 0 计算最终结果 ---
|
||
if (t_idx == 0) {{
|
||
float total_dot = s_dot[0];
|
||
float total_x_sq = s_x_sq[0];
|
||
float total_y_sq = s_y_sq[0];
|
||
|
||
float denominator = total_x_sq + total_y_sq - total_dot + eps;
|
||
|
||
output[n_idx] = (total_dot + eps) / denominator;
|
||
}}
|
||
}}
|
||
|
||
// ============================================================
|
||
// C++ Wrapper
|
||
// ============================================================
|
||
torch::Tensor tanimoto_forward_cuda(
|
||
torch::Tensor x,
|
||
torch::Tensor y,
|
||
float eps,
|
||
int N,
|
||
int D // 实例大小 (C*H*W)
|
||
) {{
|
||
// 确保输入连续
|
||
x = x.contiguous();
|
||
y = y.contiguous();
|
||
|
||
// 输出张量形状为 [N]
|
||
auto output = torch::empty({{N}}, x.options());
|
||
|
||
const int BLOCK_SIZE = 256;
|
||
dim3 threads(BLOCK_SIZE);
|
||
dim3 blocks(N);
|
||
|
||
// 动态分配共享内存 (每个累加器 256 * 4 字节)
|
||
// (dot_sum, x_sq_sum, y_sq_sum)
|
||
int shared_mem_size = 3 * BLOCK_SIZE * sizeof(float);
|
||
|
||
tanimoto_fused_kernel<<<blocks, threads, shared_mem_size>>>(
|
||
x.data_ptr<float>(),
|
||
y.data_ptr<float>(),
|
||
output.data_ptr<float>(),
|
||
eps,
|
||
N,
|
||
D
|
||
);
|
||
|
||
return output;
|
||
}}
|
||
"""
|
||
|
||
# 3. 编译 CUDA 模块
|
||
self.tanimoto_op = load_inline(
|
||
name="tanimoto_cuda_v1_fma_fixed",
|
||
cpp_sources=cpp_source,
|
||
cuda_sources=cuda_source,
|
||
functions=["tanimoto_forward_cuda"],
|
||
# (使用高精度标志, 禁用 fast_math 以确保 volatile 生效)
|
||
extra_cuda_cflags=[
|
||
"-O3",
|
||
"-ftz=false",
|
||
"-prec-div=true",
|
||
"-prec-sqrt=true",
|
||
"-std=c++17"
|
||
],
|
||
verbose=False
|
||
)
|
||
|
||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||
if x.dtype != torch.float32 or not x.is_cuda:
|
||
x = x.to("cuda", dtype=torch.float32)
|
||
if y.dtype != torch.float32 or not y.is_cuda:
|
||
y = y.to("cuda", dtype=torch.float32)
|
||
|
||
N, C, H, W = x.size()
|
||
D = C * H * W # 实例大小
|
||
|
||
return self.tanimoto_op.tanimoto_forward_cuda(
|
||
x.contiguous(),
|
||
y.contiguous(),
|
||
self.eps,
|
||
N,
|
||
D
|
||
)
|
||
|