forked from ccf-ai-infra/GPUCodeForces
132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
import math
|
|
|
|
N, C, H, W = 16, 16, 64, 64
|
|
LAMBDA = 0.5
|
|
BLOCK_SIZE = 256
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, lambd=LAMBDA):
|
|
super().__init__()
|
|
self.lambd = lambd
|
|
self.block_size = BLOCK_SIZE
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
total_elements = N * C * H * W
|
|
if total_elements % 4 != 0:
|
|
raise ValueError("Total elements must be divisible by 4 for float4 vectorization.")
|
|
|
|
lambd_str = f"{self.lambd:.8f}f"
|
|
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor softshrink_cuda_vec(torch::Tensor input, float lambd);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
#define BLOCK_SIZE {self.block_size}
|
|
|
|
/*
|
|
* Softshrink Kernel with float4 Vectorization
|
|
* Implements: y = sign(x) * max(0, |x| - lambda)
|
|
*/
|
|
__global__ void softshrink_kernel_vec(
|
|
const float* __restrict__ input,
|
|
float* __restrict__ output,
|
|
int total_vec_elements,
|
|
float lambd // lambda value passed as argument
|
|
) {{
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
for (; idx < total_vec_elements; idx += gridDim.x * blockDim.x) {{
|
|
|
|
// 1. Load float4 vector
|
|
const float4* in_ptr = reinterpret_cast<const float4*>(input);
|
|
float4 in_vec = in_ptr[idx];
|
|
|
|
// 2. Apply Softshrink element-wise
|
|
float4 out_vec;
|
|
|
|
// CUDA Intrinsics: fabsf (abs), copysignf (sign)
|
|
// Note: fabsf(x) - lambd is the core logic.
|
|
// We use copysignf to restore the sign, which is robust near zero.
|
|
|
|
// --- X component ---
|
|
float magnitude_x = fabsf(in_vec.x);
|
|
float thresholded_x = fmaxf(0.0f, magnitude_x - lambd); // max(0, |x| - lambda)
|
|
out_vec.x = copysignf(thresholded_x, in_vec.x); // sign(x) * thresholded_x
|
|
|
|
// --- Y component ---
|
|
float magnitude_y = fabsf(in_vec.y);
|
|
float thresholded_y = fmaxf(0.0f, magnitude_y - lambd);
|
|
out_vec.y = copysignf(thresholded_y, in_vec.y);
|
|
|
|
// --- Z component ---
|
|
float magnitude_z = fabsf(in_vec.z);
|
|
float thresholded_z = fmaxf(0.0f, magnitude_z - lambd);
|
|
out_vec.z = copysignf(thresholded_z, in_vec.z);
|
|
|
|
// --- W component ---
|
|
float magnitude_w = fabsf(in_vec.w);
|
|
float thresholded_w = fmaxf(0.0f, magnitude_w - lambd);
|
|
out_vec.w = copysignf(thresholded_w, in_vec.w);
|
|
|
|
|
|
// 3. Store float4 vector
|
|
float4* out_ptr = reinterpret_cast<float4*>(output);
|
|
out_ptr[idx] = out_vec;
|
|
}}
|
|
}}
|
|
|
|
|
|
torch::Tensor softshrink_cuda_vec(torch::Tensor input, float lambd) {{
|
|
|
|
TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
|
|
TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
|
|
|
|
int total_elements = input.numel();
|
|
|
|
TORCH_CHECK(total_elements % 4 == 0, "Total elements must be divisible by 4 for float4 vectorization.");
|
|
|
|
input = input.contiguous();
|
|
|
|
auto output = torch::empty_like(input);
|
|
|
|
int total_vec_elements = total_elements / 4;
|
|
|
|
int blocks = std::min((total_vec_elements + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);
|
|
|
|
softshrink_kernel_vec<<<blocks, BLOCK_SIZE>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
total_vec_elements,
|
|
lambd // Passing lambda
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = ['-O3']
|
|
|
|
self.op = load_inline(
|
|
name='softshrink_opt_vec',
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=['softshrink_cuda_vec'],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
if not input.is_cuda: input = input.cuda()
|
|
input_cont = input.contiguous()
|
|
return self.op.softshrink_cuda_vec(input_cont, self.lambd) |