forked from ccf-ai-infra/GPUCodeForces
274 lines
9.2 KiB
Python
274 lines
9.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, bins=32):
|
|
super().__init__()
|
|
self.bins = bins
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor vi_cuda(torch::Tensor x, torch::Tensor y, int bins);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <float.h>
|
|
|
|
__device__ __forceinline__ float warp_reduce_min(float val) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset /= 2) {
|
|
val = fminf(val, __shfl_down_sync(0xffffffff, val, offset));
|
|
}
|
|
return val;
|
|
}
|
|
|
|
__device__ __forceinline__ float warp_reduce_max(float val) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset /= 2) {
|
|
val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
|
|
}
|
|
return val;
|
|
}
|
|
|
|
__device__ __forceinline__ double warp_reduce_sum(double val) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset /= 2) {
|
|
val += __shfl_down_sync(0xffffffff, val, offset);
|
|
}
|
|
return val;
|
|
}
|
|
|
|
__global__ void vi_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ y,
|
|
float* __restrict__ output,
|
|
int feature_dim,
|
|
int batch_size)
|
|
{
|
|
const int BINS = 32;
|
|
|
|
__shared__ int s_hist[BINS][BINS];
|
|
__shared__ float s_px[BINS];
|
|
__shared__ float s_py[BINS];
|
|
|
|
// Min/Max reduction buffer
|
|
static __shared__ float s_reduce_buf[8][4];
|
|
|
|
int bid = blockIdx.x;
|
|
if (bid >= batch_size) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int lane = tid % 32;
|
|
int wid = tid / 32;
|
|
for (int i = tid; i < BINS * BINS; i += blockDim.x) {
|
|
reinterpret_cast<int*>(s_hist)[i] = 0;
|
|
}
|
|
if (tid < BINS) {
|
|
s_px[tid] = 0.0f;
|
|
s_py[tid] = 0.0f;
|
|
}
|
|
__syncthreads();
|
|
|
|
const float* x_row = x + bid * feature_dim;
|
|
const float* y_row = y + bid * feature_dim;
|
|
|
|
float l_min_x = FLT_MAX, l_max_x = -FLT_MAX;
|
|
float l_min_y = FLT_MAX, l_max_y = -FLT_MAX;
|
|
|
|
int vec_loops = feature_dim / 4;
|
|
int vec_remainder = feature_dim % 4;
|
|
const float4* x_vec = reinterpret_cast<const float4*>(x_row);
|
|
const float4* y_vec = reinterpret_cast<const float4*>(y_row);
|
|
|
|
for (int i = tid; i < vec_loops; i += blockDim.x) {
|
|
float4 vx = x_vec[i];
|
|
float4 vy = y_vec[i];
|
|
|
|
l_min_x = fminf(l_min_x, fminf(vx.x, fminf(vx.y, fminf(vx.z, vx.w))));
|
|
l_max_x = fmaxf(l_max_x, fmaxf(vx.x, fmaxf(vx.y, fmaxf(vx.z, vx.w))));
|
|
l_min_y = fminf(l_min_y, fminf(vy.x, fminf(vy.y, fminf(vy.z, vy.w))));
|
|
l_max_y = fmaxf(l_max_y, fmaxf(vy.x, fmaxf(vy.y, fmaxf(vy.z, vy.w))));
|
|
}
|
|
l_min_x = warp_reduce_min(l_min_x);
|
|
l_max_x = warp_reduce_max(l_max_x);
|
|
l_min_y = warp_reduce_min(l_min_y);
|
|
l_max_y = warp_reduce_max(l_max_y);
|
|
|
|
if (lane == 0) {
|
|
s_reduce_buf[wid][0] = l_min_x;
|
|
s_reduce_buf[wid][1] = l_max_x;
|
|
s_reduce_buf[wid][2] = l_min_y;
|
|
s_reduce_buf[wid][3] = l_max_y;
|
|
}
|
|
__syncthreads();
|
|
|
|
// Block Reduce (by thread 0)
|
|
float min_x, max_x, min_y, max_y;
|
|
if (tid == 0) {
|
|
min_x = FLT_MAX; max_x = -FLT_MAX;
|
|
min_y = FLT_MAX; max_y = -FLT_MAX;
|
|
for (int i = 0; i < blockDim.x / 32; ++i) {
|
|
min_x = fminf(min_x, s_reduce_buf[i][0]);
|
|
max_x = fmaxf(max_x, s_reduce_buf[i][1]);
|
|
min_y = fminf(min_y, s_reduce_buf[i][2]);
|
|
max_y = fmaxf(max_y, s_reduce_buf[i][3]);
|
|
}
|
|
// Store back to shared for broadcast
|
|
s_reduce_buf[0][0] = min_x; s_reduce_buf[0][1] = max_x;
|
|
s_reduce_buf[0][2] = min_y; s_reduce_buf[0][3] = max_y;
|
|
}
|
|
__syncthreads();
|
|
|
|
min_x = s_reduce_buf[0][0]; max_x = s_reduce_buf[0][1];
|
|
min_y = s_reduce_buf[0][2]; max_y = s_reduce_buf[0][3];
|
|
|
|
float range_x = max_x - min_x + 1e-6f;
|
|
float range_y = max_y - min_y + 1e-6f;
|
|
|
|
for (int i = tid; i < vec_loops; i += blockDim.x) {
|
|
float4 vx = x_vec[i];
|
|
float4 vy = y_vec[i];
|
|
|
|
#pragma unroll
|
|
for (int k = 0; k < 4; ++k) {
|
|
float val_x = (k==0?vx.x:k==1?vx.y:k==2?vx.z:vx.w);
|
|
float val_y = (k==0?vy.x:k==1?vy.y:k==2?vy.z:vy.w);
|
|
|
|
int bx = (int)((val_x - min_x) / range_x * BINS);
|
|
int by = (int)((val_y - min_y) / range_y * BINS);
|
|
bx = min(max(bx, 0), BINS - 1);
|
|
by = min(max(by, 0), BINS - 1);
|
|
|
|
atomicAdd(&s_hist[bx][by], 1);
|
|
}
|
|
}
|
|
|
|
if (tid == 0 && vec_remainder > 0) {
|
|
for (int i = 0; i < vec_remainder; ++i) {
|
|
int idx = vec_loops * 4 + i;
|
|
int bx = (int)((x_row[idx] - min_x) / range_x * BINS);
|
|
int by = (int)((y_row[idx] - min_y) / range_y * BINS);
|
|
bx = min(max(bx, 0), BINS - 1);
|
|
by = min(max(by, 0), BINS - 1);
|
|
atomicAdd(&s_hist[bx][by], 1);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
|
|
if (tid < BINS) {
|
|
int px = 0;
|
|
int py = 0;
|
|
for (int i = 0; i < BINS; ++i) {
|
|
px += s_hist[tid][i];
|
|
py += s_hist[i][tid];
|
|
}
|
|
s_px[tid] = (float)px / feature_dim;
|
|
s_py[tid] = (float)py / feature_dim;
|
|
}
|
|
__syncthreads();
|
|
|
|
|
|
|
|
double sum_mi = 0.0;
|
|
double sum_hx = 0.0;
|
|
double sum_hy = 0.0;
|
|
const float eps = 1e-12f;
|
|
|
|
// Compute MI (over 32x32 grid)
|
|
for (int i = tid; i < BINS * BINS; i += blockDim.x) {
|
|
int r = i / BINS;
|
|
int c = i % BINS;
|
|
|
|
int count = s_hist[r][c];
|
|
if (count > 0) {
|
|
float p_xy = (float)count / feature_dim;
|
|
float px = s_px[r];
|
|
float py = s_py[c];
|
|
|
|
sum_mi += (double)(p_xy * logf((p_xy + eps) / (px * py + eps)));
|
|
}
|
|
}
|
|
|
|
// Compute H(X) and H(Y) (over 32 bins)
|
|
// Each thread takes one bin if tid < 32
|
|
if (tid < BINS) {
|
|
float px = s_px[tid];
|
|
float py = s_py[tid];
|
|
if (px > 0) sum_hx += (double)(-px * logf(px + eps));
|
|
if (py > 0) sum_hy += (double)(-py * logf(py + eps));
|
|
}
|
|
|
|
// Reductions
|
|
sum_mi = warp_reduce_sum(sum_mi);
|
|
sum_hx = warp_reduce_sum(sum_hx);
|
|
sum_hy = warp_reduce_sum(sum_hy);
|
|
|
|
static __shared__ double s_vals[3][32]; // [0]:MI, [1]:HX, [2]:HY
|
|
|
|
if (lane == 0) {
|
|
s_vals[0][wid] = sum_mi;
|
|
s_vals[1][wid] = sum_hx;
|
|
s_vals[2][wid] = sum_hy;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (tid == 0) {
|
|
double total_mi = 0.0, total_hx = 0.0, total_hy = 0.0;
|
|
int n_warps = blockDim.x / 32;
|
|
|
|
for (int i = 0; i < n_warps; ++i) {
|
|
total_mi += s_vals[0][i];
|
|
total_hx += s_vals[1][i];
|
|
total_hy += s_vals[2][i];
|
|
}
|
|
|
|
// VI = H(X) + H(Y) - 2 * MI
|
|
output[bid] = (float)(total_hx + total_hy - 2.0 * total_mi);
|
|
}
|
|
}
|
|
|
|
torch::Tensor vi_cuda(torch::Tensor x, torch::Tensor y, int bins) {
|
|
auto x_c = x.contiguous();
|
|
auto y_c = y.contiguous();
|
|
|
|
int batch_size = x_c.size(0);
|
|
int feature_dim = x_c.size(1);
|
|
|
|
auto output = torch::empty({batch_size}, x.options());
|
|
|
|
int threads = 256;
|
|
int blocks = batch_size;
|
|
|
|
if (bins != 32) {
|
|
// Fallback logic could be added here
|
|
}
|
|
|
|
vi_kernel<<<blocks, threads>>>(
|
|
x_c.data_ptr<float>(),
|
|
y_c.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
feature_dim,
|
|
batch_size
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name="vi_opt_v1",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["vi_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, x, y):
|
|
return self.op.vi_cuda(x, y, self.bins) |