forked from ccf-ai-infra/GPUCodeForces
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
N, C, H, W = 16, 16, 64, 64
|
|
BETA = 1.0
|
|
THRESHOLD = 20.0
|
|
BLOCK_SIZE = 256
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, beta=BETA, threshold=THRESHOLD):
|
|
super().__init__()
|
|
self.beta = beta
|
|
self.threshold = threshold
|
|
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.")
|
|
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor softplus_cuda_vec(torch::Tensor input, float beta, float threshold);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
#define BLOCK_SIZE {self.block_size}
|
|
|
|
__device__ __forceinline__ float softplus_elem(float x, float beta, float threshold, float inv_beta) {{
|
|
float scaled_x = x * beta;
|
|
return (scaled_x > threshold) ? x : inv_beta * __logf(1.0f + __expf(scaled_x));
|
|
}}
|
|
|
|
__global__ void softplus_kernel_vec(
|
|
const float4* __restrict__ input,
|
|
float4* __restrict__ output,
|
|
int total_vec_elements,
|
|
float beta,
|
|
float threshold,
|
|
float inv_beta
|
|
) {{
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int stride = gridDim.x * blockDim.x;
|
|
|
|
for (int i = idx; i < total_vec_elements; i += stride) {{
|
|
float4 in_vec = input[i];
|
|
float4 out_vec;
|
|
|
|
out_vec.x = softplus_elem(in_vec.x, beta, threshold, inv_beta);
|
|
out_vec.y = softplus_elem(in_vec.y, beta, threshold, inv_beta);
|
|
out_vec.z = softplus_elem(in_vec.z, beta, threshold, inv_beta);
|
|
out_vec.w = softplus_elem(in_vec.w, beta, threshold, inv_beta);
|
|
|
|
output[i] = out_vec;
|
|
}}
|
|
}}
|
|
|
|
torch::Tensor softplus_cuda_vec(torch::Tensor input, float beta, float threshold) {{
|
|
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;
|
|
float inv_beta = 1.0f / beta;
|
|
|
|
int blocks = std::min((total_vec_elements + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);
|
|
|
|
softplus_kernel_vec<<<blocks, BLOCK_SIZE>>>(
|
|
reinterpret_cast<const float4*>(input.data_ptr<float>()),
|
|
reinterpret_cast<float4*>(output.data_ptr<float>()),
|
|
total_vec_elements,
|
|
beta,
|
|
threshold,
|
|
inv_beta
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = ['-O3', '--use_fast_math']
|
|
|
|
self.op = load_inline(
|
|
name='softplus_opt_vec',
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=['softplus_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.softplus_cuda_vec(input_cont, self.beta, self.threshold) |