forked from ccf-ai-infra/GPUCodeForces
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__global__ void affine_relu_kernel(const float* x, const float* scale, const float* bias, float* y, int dim, long long total) {
|
|
long long idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
long long stride = blockDim.x * gridDim.x;
|
|
for (long long i = idx; i < total; i += stride) {
|
|
int j = (int)(i % dim);
|
|
float v = x[i] * scale[j] + bias[j];
|
|
y[i] = v > 0.f ? v : 0.f;
|
|
}
|
|
}
|
|
|
|
torch::Tensor affine_relu_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias) {
|
|
auto x_contig = x.contiguous();
|
|
auto s_contig = scale.contiguous();
|
|
auto b_contig = bias.contiguous();
|
|
auto y = torch::empty_like(x_contig);
|
|
long long total = x_contig.numel();
|
|
int dim = (int)x_contig.size(-1);
|
|
int block = 512;
|
|
long long grid = (total + block - 1) / block;
|
|
grid = grid > 65535 ? 65535 : grid;
|
|
affine_relu_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), s_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total);
|
|
return y;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor affine_relu_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias);
|
|
"""
|
|
|
|
ops = load_inline(
|
|
name="affine_relu",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=source,
|
|
functions=["affine_relu_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self, scale: torch.Tensor, bias: torch.Tensor):
|
|
super(ModelNew, self).__init__()
|
|
self.ops = ops
|
|
self.register_buffer("scale", scale)
|
|
self.register_buffer("bias", bias)
|
|
|
|
def forward(self, x):
|
|
return self.ops.affine_relu_cuda(x, self.scale, self.bias)
|