forked from ccf-ai-infra/GPUCodeForces
161 lines
6.2 KiB
Python
161 lines
6.2 KiB
Python
# switchablenorm_cuda.py
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from switchablenorm_torch import N, C, H, W, EPS
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, weight, bias, w_in, w_ln, w_bn):
|
|
super().__init__()
|
|
self.register_buffer('weight', weight)
|
|
self.register_buffer('bias', bias)
|
|
self.register_buffer('w_in', w_in)
|
|
self.register_buffer('w_ln', w_ln)
|
|
self.register_buffer('w_bn', w_bn)
|
|
self.eps = EPS
|
|
self.register_buffer('running_mean', torch.zeros(C))
|
|
self.register_buffer('running_var', torch.ones(C))
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
|
|
torch::Tensor sn_forward_cuda(
|
|
torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
|
|
torch::Tensor w_in, torch::Tensor w_ln, torch::Tensor w_bn,
|
|
torch::Tensor running_mean, torch::Tensor running_var,
|
|
float eps, int N, int C, int H, int W);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <torch/extension.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
#define WARP_SIZE 32
|
|
|
|
__device__ __forceinline__ float warp_reduce_sum(float val) {{
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
|
|
val += __shfl_down_sync(0xffffffff, val, offset);
|
|
return val;
|
|
}}
|
|
|
|
__global__ void sn_normalize_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ ln_mean_ptr, const float* __restrict__ ln_var_ptr,
|
|
const float* __restrict__ in_mean_ptr, const float* __restrict__ in_var_ptr,
|
|
const float* __restrict__ bn_mean_ptr, const float* __restrict__ bn_var_ptr,
|
|
const float* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
|
|
const float* __restrict__ w_in_ptr, const float* __restrict__ w_ln_ptr, const float* __restrict__ w_bn_ptr,
|
|
float* __restrict__ output,
|
|
int N, int C, int H, int W, float eps
|
|
) {{
|
|
int n_elements = N * C * H * W;
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (idx >= n_elements) return;
|
|
|
|
int nc_idx = idx / (H * W);
|
|
int sample_idx = nc_idx / C;
|
|
int channel_idx = nc_idx % C;
|
|
|
|
float w_in = w_in_ptr[channel_idx];
|
|
float w_ln = w_ln_ptr[channel_idx];
|
|
float w_bn = w_bn_ptr[channel_idx];
|
|
|
|
float w_sum = fabsf(w_in) + fabsf(w_ln) + fabsf(w_bn);
|
|
float w_in_norm = fabsf(w_in) / w_sum;
|
|
float w_ln_norm = fabsf(w_ln) / w_sum;
|
|
float w_bn_norm = fabsf(w_bn) / w_sum;
|
|
|
|
float mean_ln = ln_mean_ptr[sample_idx];
|
|
float var_ln = ln_var_ptr[sample_idx];
|
|
|
|
float mean_in = in_mean_ptr[nc_idx];
|
|
float var_in = in_var_ptr[nc_idx];
|
|
|
|
float mean_bn = bn_mean_ptr[channel_idx];
|
|
float var_bn = bn_var_ptr[channel_idx];
|
|
|
|
float agg_mean = w_in_norm * mean_in + w_ln_norm * mean_ln + w_bn_norm * mean_bn;
|
|
|
|
float m2_in = var_in + mean_in * mean_in;
|
|
float m2_ln = var_ln + mean_ln * mean_ln;
|
|
float m2_bn = var_bn + mean_bn * mean_bn;
|
|
|
|
float agg_m2 = w_in_norm * m2_in + w_ln_norm * m2_ln + w_bn_norm * m2_bn;
|
|
float agg_var = agg_m2 - agg_mean * agg_mean;
|
|
|
|
float val = x[idx];
|
|
float gamma = weight_ptr[channel_idx];
|
|
float beta = bias_ptr[channel_idx];
|
|
|
|
float inv_std = rsqrtf(agg_var + eps);
|
|
float normalized = (val - agg_mean) * inv_std;
|
|
|
|
output[idx] = normalized * gamma + beta;
|
|
}}
|
|
|
|
|
|
torch::Tensor sn_forward_cuda(
|
|
torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
|
|
torch::Tensor w_in, torch::Tensor w_ln, torch::Tensor w_bn,
|
|
torch::Tensor running_mean, torch::Tensor running_var,
|
|
float eps, int N, int C, int H, int W) {{
|
|
|
|
|
|
torch::Tensor in_mean = input.mean({{2, 3}}, true).squeeze(-1).squeeze(-1).contiguous();
|
|
torch::Tensor in_var = input.var({{2, 3}}, true).squeeze(-1).squeeze(-1).contiguous();
|
|
|
|
torch::Tensor ln_mean = input.mean({{1, 2, 3}}, true).squeeze(-1).squeeze(-1).squeeze(-1).contiguous();
|
|
torch::Tensor ln_var = input.var({{1, 2, 3}}, true).squeeze(-1).squeeze(-1).squeeze(-1).contiguous();
|
|
|
|
auto output = torch::empty_like(input).contiguous();
|
|
|
|
int n_elements = input.numel();
|
|
const int blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
|
|
|
sn_normalize_kernel<<<blocks, BLOCK_SIZE>>>(
|
|
input.data_ptr<float>(),
|
|
ln_mean.data_ptr<float>(), ln_var.data_ptr<float>(),
|
|
in_mean.data_ptr<float>(), in_var.data_ptr<float>(),
|
|
running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
|
|
weight.data_ptr<float>(), bias.data_ptr<float>(),
|
|
w_in.data_ptr<float>(), w_ln.data_ptr<float>(), w_bn.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
N, C, H, W, eps
|
|
);
|
|
|
|
return output;
|
|
}}
|
|
"""
|
|
|
|
self.sn_op = load_inline(
|
|
name="sn_fused_op_logic_correct_v3",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["sn_forward_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=True
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
weight = self.weight.squeeze().contiguous()
|
|
bias = self.bias.squeeze().contiguous()
|
|
w_in = self.w_in.contiguous()
|
|
w_ln = self.w_ln.contiguous()
|
|
w_bn = self.w_bn.contiguous()
|
|
running_mean = self.running_mean.contiguous()
|
|
running_var = self.running_var.contiguous()
|
|
|
|
N, C, H, W = x.size(0), x.size(1), x.size(2), x.size(3)
|
|
|
|
return self.sn_op.sn_forward_cuda(
|
|
x.contiguous(), weight, bias, w_in, w_ln, w_bn,
|
|
running_mean, running_var, self.eps, N, C, H, W
|
|
) |