forked from ccf-ai-infra/GPUCodeForces
98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
# localresponsenorm_cuda.py
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from localresponsenorm_torch import N, C, H, W, SIZE, ALPHA, BETA, K
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
half_size = SIZE // 2
|
|
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <device_launch_parameters.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
|
|
__global__ void lrn_fused_kernel(
|
|
const float* __restrict__ x,
|
|
float* __restrict__ y,
|
|
int N, int C, int H, int W,
|
|
int size, int half_size, float alpha, float beta, float k
|
|
) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx >= N * C * H * W) return;
|
|
|
|
int n = idx / (C * H * W);
|
|
int rem = idx % (C * H * W);
|
|
int c = rem / (H * W);
|
|
int hw = rem % (H * W);
|
|
|
|
int base_idx = n * C * H * W + hw;
|
|
float sum_sq = 0.0f;
|
|
|
|
int start_c = max(0, c - half_size);
|
|
int end_c = min(C - 1, c + half_size);
|
|
|
|
#pragma unroll 4
|
|
for (int window_c = start_c; window_c <= end_c; ++window_c) {
|
|
float val = x[base_idx + window_c * H * W];
|
|
sum_sq = __fmaf_rn(val, val, sum_sq);
|
|
}
|
|
|
|
float alpha_over_n = alpha / (float)size;
|
|
float scale = __fmaf_rn(alpha_over_n, sum_sq, k);
|
|
|
|
|
|
float norm_factor = __powf(scale, beta);
|
|
float input_val = x[idx];
|
|
y[idx] = __fdividef(input_val, norm_factor);
|
|
}
|
|
|
|
torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k) {
|
|
const int N = input.size(0);
|
|
const int C = input.size(1);
|
|
const int H = input.size(2);
|
|
const int W = input.size(3);
|
|
auto output = torch::empty_like(input);
|
|
|
|
const int n_elements = N * C * H * W;
|
|
const int blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
|
|
|
lrn_fused_kernel<<<blocks, BLOCK_SIZE>>>(
|
|
input.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
N, C, H, W,
|
|
size, half_size, alpha, beta, k
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
self.lrn_op = load_inline(
|
|
name="lrn_correct_formula",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["lrn_forward_cuda"],
|
|
extra_cuda_cflags=[
|
|
"-O3",
|
|
"--use_fast_math",
|
|
"--fmad=true"
|
|
],
|
|
verbose=True
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
half_size = SIZE // 2
|
|
return self.lrn_op.lrn_forward_cuda(x.contiguous(), SIZE, half_size, ALPHA, BETA, K) |