
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void simple_lrn_kernel(
    const float* __restrict__ x, 
    float* __restrict__ y, 
    int N, int C, int H, int W,
    float alpha, float beta, float k
) {
    int w = threadIdx.x;
    int h = blockIdx.x;
    int c = blockIdx.y;
    int n = blockIdx.z;
    
    if (h >= H || w >= W || c >= C) return;
    
    int idx = n * C * H * W + c * H * W + h * W + w;
    float val = x[idx];
    
    // 计算相邻通道的平方和（匹配size=3的LRN）
    float sum_sq = 0.0f;
    for (int j = max(0, c-1); j <= min(C-1, c+1); ++j) {
        int j_idx = n * C * H * W + j * H * W + h * W + w;
        sum_sq += x[j_idx] * x[j_idx];
    }
    
    // 归一化计算
    y[idx] = val / powf(k + alpha * sum_sq, beta);
}

torch::Tensor simple_lrn_cuda(torch::Tensor x) {
    x = x.contiguous();  
    auto dims = x.sizes(); 
    int N = dims[0], C = dims[1], H = dims[2], W = dims[3];
    auto y = torch::empty_like(x);
    
    // 线程配置：W为线程数，H×C×N为网格数
    dim3 block(W);
    dim3 grid(H, C, N);
    
    simple_lrn_kernel<<<grid, block>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        N, C, H, W,
        1e-4f, 0.75f, 2.0f  
    );
    return y;
}
"""

cpp_source = "torch::Tensor simple_lrn_cuda(torch::Tensor x);"

# 编译CUDA代码
simple_lrn = load_inline(
    name="simple_lrn",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["simple_lrn_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)

class ModelNew(nn.Module):  
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return simple_lrn.simple_lrn_cuda(x)


def get_init_inputs():
    return []


def get_inputs():
    return [torch.randn(8, 32, 64, 64)]  