import torch
from torch.utils.cpp_extension import load_inline

source = """
#include <torch/extension.h>

__inline__ __device__ float warpReduceSum(float val){
    for(int offset=16; offset>0; offset>>=1){
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void channel_rms_gate_kernel(const float* x, const float* gamma, const float* beta, float* y, int N, int C, int H, int W){
    int b = blockIdx.x / C;
    int c = blockIdx.x % C;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warpId = tid >> 5;
    extern __shared__ float sh[];
    int HW = H * W;
    float sumsq = 0.0f;
    for(int i = tid; i < HW; i += blockDim.x){
        long long idx = ((long long)b * C + c) * (long long)H * (long long)W + i;
        float v = x[idx];
        sumsq += v * v;
    }
    sumsq = warpReduceSum(sumsq);
    if(lane == 0) sh[warpId] = sumsq;
    __syncthreads();
    if(warpId == 0){
        float v = (lane < ((blockDim.x + 31) >> 5)) ? sh[lane] : 0.0f;
        v = warpReduceSum(v);
        if(lane == 0) sh[0] = v;
    }
    __syncthreads();
    float rms = sqrtf(sh[0] / (float)HW);
    float g = 1.0f / (1.0f + expf(-(gamma[c] * rms + beta[c])));
    for(int i = tid; i < HW; i += blockDim.x){
        long long idx = ((long long)b * C + c) * (long long)H * (long long)W + i;
        y[idx] = x[idx] * g;
    }
}

torch::Tensor channel_rms_gate_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta){
    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto y = torch::empty_like(xc);
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    int block = 256;
    int grid = N * C;
    size_t shmem = (size_t)((block + 31) >> 5) * sizeof(float);
    channel_rms_gate_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}
"""

cpp_source = """
torch::Tensor channel_rms_gate_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta);
"""

ops = load_inline(
    name="channel_rms_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["channel_rms_gate_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, gamma: torch.Tensor, beta: torch.Tensor):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("gamma", gamma)
        self.register_buffer("beta", beta)

    def forward(self, x: torch.Tensor):
        return self.ops.channel_rms_gate_cuda(x, self.gamma, self.beta)
