forked from ccf-ai-infra/GPUCodeForces
finish hamming_relu #140
This commit is contained in:
parent
073bc63c13
commit
e4be835ddd
|
|
@ -1,228 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
__inline__ __device__ float warpReduceMax(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float blockReduceMax(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
int num_warps = blockDim.x / 32;
|
||||
|
||||
val = warpReduceMax(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
float res = (threadIdx.x < num_warps) ? shared[threadIdx.x] : -1e38f;
|
||||
if (threadIdx.x < 32) res = warpReduceMax(res);
|
||||
return res;
|
||||
}
|
||||
|
||||
__inline__ __device__ float warpReduceSum(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float blockReduceSum(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
int num_warps = blockDim.x / 32;
|
||||
|
||||
val = warpReduceSum(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
float res = (threadIdx.x < num_warps) ? shared[threadIdx.x] : 0.0f;
|
||||
if (threadIdx.x < 32) res = warpReduceSum(res);
|
||||
return res;
|
||||
}
|
||||
|
||||
__global__ void normalize_kernel(const float* __restrict__ input,
|
||||
float* output,
|
||||
int rows,
|
||||
int cols) {
|
||||
int bid = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
if (bid >= rows) return;
|
||||
|
||||
const float* row_in = input + bid * cols;
|
||||
float* row_out = output + bid * cols;
|
||||
|
||||
float sum_sq = 0.0f;
|
||||
for (int i = tid; i < cols; i += blockDim.x) {
|
||||
float val = row_in[i];
|
||||
sum_sq += val * val;
|
||||
}
|
||||
|
||||
sum_sq = blockReduceSum(sum_sq);
|
||||
|
||||
__shared__ float inv_norm;
|
||||
if (tid == 0) {
|
||||
inv_norm = rsqrtf(sum_sq + 1e-12f);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float scale = inv_norm;
|
||||
for (int i = tid; i < cols; i += blockDim.x) {
|
||||
row_out[i] = row_in[i] * scale;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void swav_loss_kernel(
|
||||
const float* __restrict__ dots1,
|
||||
const float* __restrict__ dots2,
|
||||
float* loss,
|
||||
int batch_size,
|
||||
int num_proto,
|
||||
float temp,
|
||||
float eps
|
||||
) {
|
||||
int bid = blockIdx.x;
|
||||
if (bid >= batch_size) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* row1 = dots1 + bid * num_proto;
|
||||
const float* row2 = dots2 + bid * num_proto;
|
||||
|
||||
float local_max1 = -1e38f;
|
||||
float local_max2 = -1e38f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
local_max1 = fmaxf(local_max1, row1[i]);
|
||||
local_max2 = fmaxf(local_max2, row2[i]);
|
||||
}
|
||||
|
||||
float max_dot1 = blockReduceMax(local_max1);
|
||||
float max_dot2 = blockReduceMax(local_max2);
|
||||
|
||||
__shared__ float s_max1, s_max2;
|
||||
if (tid == 0) { s_max1 = max_dot1; s_max2 = max_dot2; }
|
||||
__syncthreads();
|
||||
max_dot1 = s_max1;
|
||||
max_dot2 = s_max2;
|
||||
|
||||
float inv_temp = 1.0f / temp;
|
||||
float inv_temp_eps = 1.0f / (temp * eps);
|
||||
|
||||
float l_sum_p1 = 0.0f, l_sum_p2 = 0.0f;
|
||||
float l_sum_q1 = 0.0f, l_sum_q2 = 0.0f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
float d1 = row1[i];
|
||||
float d2 = row2[i];
|
||||
|
||||
l_sum_p1 += expf((d1 - max_dot1) * inv_temp);
|
||||
l_sum_p2 += expf((d2 - max_dot2) * inv_temp);
|
||||
|
||||
l_sum_q1 += expf((d1 - max_dot1) * inv_temp_eps);
|
||||
l_sum_q2 += expf((d2 - max_dot2) * inv_temp_eps);
|
||||
}
|
||||
|
||||
float sum_p1 = blockReduceSum(l_sum_p1);
|
||||
float sum_p2 = blockReduceSum(l_sum_p2);
|
||||
float sum_q1 = blockReduceSum(l_sum_q1);
|
||||
float sum_q2 = blockReduceSum(l_sum_q2);
|
||||
|
||||
__shared__ float s_log_sum_p1, s_log_sum_p2, s_sum_q1, s_sum_q2;
|
||||
if (tid == 0) {
|
||||
s_log_sum_p1 = logf(sum_p1);
|
||||
s_log_sum_p2 = logf(sum_p2);
|
||||
s_sum_q1 = sum_q1;
|
||||
s_sum_q2 = sum_q2;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float l_loss = 0.0f;
|
||||
|
||||
for (int i = tid; i < num_proto; i += blockDim.x) {
|
||||
float d1 = row1[i];
|
||||
float d2 = row2[i];
|
||||
|
||||
float q1 = expf((d1 - max_dot1) * inv_temp_eps) / s_sum_q1;
|
||||
float q2 = expf((d2 - max_dot2) * inv_temp_eps) / s_sum_q2;
|
||||
|
||||
float log_p1 = (d1 - max_dot1) * inv_temp - s_log_sum_p1;
|
||||
float log_p2 = (d2 - max_dot2) * inv_temp - s_log_sum_p2;
|
||||
|
||||
l_loss += q1 * log_p2 + q2 * log_p1;
|
||||
}
|
||||
|
||||
float block_loss = blockReduceSum(l_loss);
|
||||
|
||||
if (tid == 0) {
|
||||
atomicAdd(loss, -0.5f * block_loss / batch_size);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor swav_forward_cuda(torch::Tensor z1, torch::Tensor z2, torch::Tensor prototypes, float temperature, float epsilon) {
|
||||
auto z1_c = z1.contiguous();
|
||||
auto z2_c = z2.contiguous();
|
||||
auto p_c = prototypes.contiguous();
|
||||
|
||||
int batch_size = z1.size(0);
|
||||
int dim = z1.size(1);
|
||||
int num_proto = prototypes.size(0);
|
||||
|
||||
auto z1_n = torch::empty_like(z1_c);
|
||||
auto z2_n = torch::empty_like(z2_c);
|
||||
auto p_n = torch::empty_like(p_c);
|
||||
|
||||
int norm_block = 128;
|
||||
normalize_kernel<<<batch_size, norm_block>>>(z1_c.data_ptr<float>(), z1_n.data_ptr<float>(), batch_size, dim);
|
||||
normalize_kernel<<<batch_size, norm_block>>>(z2_c.data_ptr<float>(), z2_n.data_ptr<float>(), batch_size, dim);
|
||||
normalize_kernel<<<num_proto, norm_block>>>(p_c.data_ptr<float>(), p_n.data_ptr<float>(), num_proto, dim);
|
||||
|
||||
auto dots1 = torch::matmul(z1_n, p_n.transpose(0, 1));
|
||||
auto dots2 = torch::matmul(z2_n, p_n.transpose(0, 1));
|
||||
|
||||
auto loss = torch::zeros({1}, z1.options());
|
||||
|
||||
swav_loss_kernel<<<batch_size, 256>>>(
|
||||
dots1.data_ptr<float>(),
|
||||
dots2.data_ptr<float>(),
|
||||
loss.data_ptr<float>(),
|
||||
batch_size,
|
||||
num_proto,
|
||||
temperature,
|
||||
epsilon
|
||||
);
|
||||
|
||||
return loss;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor swav_forward_cuda(torch::Tensor z1, torch::Tensor z2, torch::Tensor prototypes, float temperature, float epsilon);
|
||||
"""
|
||||
|
||||
swav_module = load_inline(
|
||||
name="swav_loss_opt_v2",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["swav_forward_cuda"],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
super(ModelNew, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, z1, z2, prototypes):
|
||||
return swav_module.swav_forward_cuda(z1, z2, prototypes, self.temperature, self.epsilon)
|
||||
|
|
@ -1,50 +0,0 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, z1: torch.Tensor, z2: torch.Tensor, prototypes: torch.Tensor) -> torch.Tensor:
|
||||
z1 = torch.nn.functional.normalize(z1, dim=1)
|
||||
z2 = torch.nn.functional.normalize(z2, dim=1)
|
||||
prototypes = torch.nn.functional.normalize(prototypes, dim=1)
|
||||
|
||||
scores_1 = torch.matmul(z1, prototypes.T) / self.temperature
|
||||
scores_2 = torch.matmul(z2, prototypes.T) / self.temperature
|
||||
|
||||
with torch.no_grad():
|
||||
q1 = torch.exp(scores_1 / self.epsilon)
|
||||
q1 = q1 / q1.sum(dim=1, keepdim=True)
|
||||
|
||||
q2 = torch.exp(scores_2 / self.epsilon)
|
||||
q2 = q2 / q2.sum(dim=1, keepdim=True)
|
||||
|
||||
p1 = torch.nn.functional.softmax(scores_1, dim=1)
|
||||
p2 = torch.nn.functional.softmax(scores_2, dim=1)
|
||||
|
||||
loss = -0.5 * (torch.mean(torch.sum(q1 * torch.log(p2 + 1e-8), dim=1)) +
|
||||
torch.mean(torch.sum(q2 * torch.log(p1 + 1e-8), dim=1)))
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
batch_size = 16
|
||||
dim = 128
|
||||
num_prototypes = 3000
|
||||
|
||||
|
||||
def get_inputs():
|
||||
z1 = torch.randn(batch_size, dim)
|
||||
z2 = torch.randn(batch_size, dim)
|
||||
prototypes = torch.randn(num_prototypes, dim)
|
||||
return [z1, z2, prototypes]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
temperature = 0.1
|
||||
epsilon = 0.05
|
||||
return [temperature, epsilon]
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
__inline__ __device__ float warp_reduce(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void hamming_relu_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ y,
|
||||
int batch_size,
|
||||
int width)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
if (row >= batch_size) return;
|
||||
|
||||
const float* row_x = x + row * width;
|
||||
|
||||
float sum_abs = 0.0f;
|
||||
|
||||
for (int i = tid; i < width; i += blockDim.x) {
|
||||
float val = row_x[i];
|
||||
float t = target[i];
|
||||
sum_abs += fabsf(val - t);
|
||||
}
|
||||
|
||||
sum_abs = warp_reduce(sum_abs);
|
||||
|
||||
static __shared__ float shared_mem[32];
|
||||
int lane = tid % 32;
|
||||
int wid = tid / 32;
|
||||
|
||||
if (lane == 0) shared_mem[wid] = sum_abs;
|
||||
__syncthreads();
|
||||
|
||||
sum_abs = (tid < blockDim.x / 32) ? shared_mem[lane] : 0.0f;
|
||||
if (wid == 0) sum_abs = warp_reduce(sum_abs);
|
||||
|
||||
if (tid == 0) {
|
||||
y[row] = fmaxf(sum_abs, 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor launch_hamming_relu(torch::Tensor x, torch::Tensor target) {
|
||||
auto batch_size = x.size(0);
|
||||
auto width = x.size(1);
|
||||
auto y = torch::empty({batch_size}, x.options());
|
||||
|
||||
const int threads = 256;
|
||||
const int blocks = batch_size;
|
||||
|
||||
hamming_relu_kernel<<<blocks, threads>>>(
|
||||
x.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch_size,
|
||||
width
|
||||
);
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor launch_hamming_relu(torch::Tensor x, torch::Tensor target);
|
||||
"""
|
||||
|
||||
hamming_relu_module = load_inline(
|
||||
name='hamming_relu_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['launch_hamming_relu'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, target):
|
||||
super(ModelNew, self).__init__()
|
||||
self.target = nn.Parameter(target)
|
||||
self.op = hamming_relu_module
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.launch_hamming_relu(x.contiguous(), self.target.contiguous())
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, target):
|
||||
super(Model, self).__init__()
|
||||
self.target = nn.Parameter(target)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
dist = torch.sum(torch.abs(x - self.target), dim=-1)
|
||||
return torch.relu(dist)
|
||||
|
||||
batch_size = 128
|
||||
input_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, input_dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
target = torch.randn(input_dim)
|
||||
return [target]
|
||||
|
|
@ -1,27 +1,19 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
Custom CUDA kernel extension via torch.utils.cpp_extension.load_inline
|
||||
CUDA C++ kernel for L1 (Hamming‑like) distance with ReLU activation
|
||||
|
||||
SwAV (Swapping Assignments between Views) loss computation
|
||||
Element‑wise absolute differences: |x[i] – target[i]| accumulated across features
|
||||
|
||||
Multi-kernel design: normalization + loss computation
|
||||
Two‑level parallel reduction: warp‑level (__shfl_down_sync) + shared‑memory reduction
|
||||
|
||||
Warp-level reduction utilities for max/sum operations
|
||||
ReLU activation: max(distance, 0) applied after reduction (distance is non‑negative, so identity)
|
||||
|
||||
Online clustering with Sinkhorn-Knopp approximation (via epsilon scaling)
|
||||
Grid‑stride loops for coalesced memory access across feature dimension
|
||||
|
||||
Prototype-based contrastive learning with temperature scaling
|
||||
Block‑per‑sample processing with 256 threads per block
|
||||
|
||||
Row-wise L2 normalization with shared memory optimization
|
||||
|
||||
Matrix multiplication for prototype assignment scores
|
||||
|
||||
Numerically stable softmax with max subtraction
|
||||
|
||||
Atomic addition (atomicAdd) for loss accumulation
|
||||
|
||||
Symmetric loss computation between two augmented views
|
||||
PyTorch inline C++/CUDA extension via load_inline
|
||||
|
||||
|
||||
|
||||
|
|
@ -30,50 +22,22 @@ Here's an example to show you the syntax of inline embedding custom CUDA operato
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, temperature, epsilon):
|
||||
def __init__(self, target):
|
||||
super(Model, self).__init__()
|
||||
self.temperature = temperature
|
||||
self.epsilon = epsilon
|
||||
self.target = nn.Parameter(target)
|
||||
|
||||
def forward(self, z1: torch.Tensor, z2: torch.Tensor, prototypes: torch.Tensor) -> torch.Tensor:
|
||||
z1 = torch.nn.functional.normalize(z1, dim=1)
|
||||
z2 = torch.nn.functional.normalize(z2, dim=1)
|
||||
prototypes = torch.nn.functional.normalize(prototypes, dim=1)
|
||||
|
||||
scores_1 = torch.matmul(z1, prototypes.T) / self.temperature
|
||||
scores_2 = torch.matmul(z2, prototypes.T) / self.temperature
|
||||
|
||||
with torch.no_grad():
|
||||
q1 = torch.exp(scores_1 / self.epsilon)
|
||||
q1 = q1 / q1.sum(dim=1, keepdim=True)
|
||||
|
||||
q2 = torch.exp(scores_2 / self.epsilon)
|
||||
q2 = q2 / q2.sum(dim=1, keepdim=True)
|
||||
|
||||
p1 = torch.nn.functional.softmax(scores_1, dim=1)
|
||||
p2 = torch.nn.functional.softmax(scores_2, dim=1)
|
||||
|
||||
loss = -0.5 * (torch.mean(torch.sum(q1 * torch.log(p2 + 1e-8), dim=1)) +
|
||||
torch.mean(torch.sum(q2 * torch.log(p1 + 1e-8), dim=1)))
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
batch_size = 16
|
||||
dim = 128
|
||||
num_prototypes = 3000
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
dist = torch.sum(torch.abs(x - self.target), dim=-1)
|
||||
return torch.relu(dist)
|
||||
|
||||
batch_size = 128
|
||||
input_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
z1 = torch.randn(batch_size, dim)
|
||||
z2 = torch.randn(batch_size, dim)
|
||||
prototypes = torch.randn(num_prototypes, dim)
|
||||
return [z1, z2, prototypes]
|
||||
|
||||
x = torch.randn(batch_size, input_dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
temperature = 0.1
|
||||
epsilon = 0.05
|
||||
return [temperature, epsilon]
|
||||
target = torch.randn(input_dim)
|
||||
return [targe
|
||||
|
|
@ -4,8 +4,8 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from SwAVLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from SwAVLoss_cuda import ModelNew
|
||||
from hamming_relu_torch import Model, get_inputs, get_init_inputs
|
||||
from hamming_relu_cuda import ModelNew
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
|
|
|
|||
Loading…
Reference in New Issue