import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Pre-norm 1D dilated convolution block with BatchNorm: BatchNorm1d, ReLU activation, dilated Conv1d, and residual addition.
    At high dilation rates, memory access becomes non-contiguous across large strides in the sequence dimension.
    """
    def __init__(self, channels, kernel_size, dilation):
        super(Model, self).__init__()
        self.norm = nn.BatchNorm1d(channels)
        self.activation = nn.ReLU()
        self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)

    def forward(self, x):
        residual = x
        x = self.norm(x)
        x = self.activation(x)
        x = self.conv(x)
        x = x + residual
        return x

batch_size = 32
channels = 256
seq_len = 4096
kernel_size = 3
dilation = 128

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), channels, popcorn_pri.jitter_int(seq_len))]

def get_init_inputs():
    return [channels, kernel_size, dilation]

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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor bn_relu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                           torch::Tensor running_mean, torch::Tensor running_var,
                           bool training, double momentum, double eps);

torch::Tensor add_residual_cuda(torch::Tensor y, torch::Tensor x);

torch::Tensor bn_relu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              bool training, double momentum, double eps) {
    return bn_relu_cuda(x, weight, bias, running_mean, running_var, training, momentum, eps);
}

torch::Tensor add_residual(torch::Tensor y, torch::Tensor x) {
    return add_residual_cuda(y, x);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define BN_THREADS 256

__global__ void bn_relu_train_kernel(const float* __restrict__ x,
                                     float* __restrict__ out,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ running_mean,
                                     float* __restrict__ running_var,
                                     int N, int C, int L,
                                     float momentum, float eps) {
    int c = blockIdx.x;
    int tid = threadIdx.x;
    int M = N * L;
    __shared__ float ssum[BN_THREADS];
    __shared__ float ssq[BN_THREADS];
    __shared__ float smu;
    __shared__ float sinv;
    float sum = 0.0f;
    float sumsq = 0.0f;

    for (int i = tid; i < M; i += BN_THREADS) {
        int n = i / L;
        int l = i - n * L;
        float v = x[(n * C + c) * L + l];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    ssum[tid] = sum;
    ssq[tid] = sumsq;
    __syncthreads();

    for (int stride = BN_THREADS >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            ssq[tid] += ssq[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        float mean = ssum[0] / (float)M;
        float var = ssq[0] / (float)M - mean * mean;
        var = var < 0.0f ? 0.0f : var;
        smu = mean;
        sinv = rsqrtf(var + eps);
        // PyTorch BatchNorm updates running_var with the unbiased estimator.
        if (running_mean != nullptr && running_var != nullptr) {
            float unbiased = var * ((float)M / (float)(M - 1));
            running_mean[c] = running_mean[c] * (1.0f - momentum) + mean * momentum;
            running_var[c]  = running_var[c]  * (1.0f - momentum) + unbiased * momentum;
        }
    }
    __syncthreads();

    float g = gamma[c];
    float b = beta[c];
    float mean = smu;
    float inv = sinv;
    for (int i = tid; i < M; i += BN_THREADS) {
        int n = i / L;
        int l = i - n * L;
        int idx = (n * C + c) * L + l;
        float v = (x[idx] - mean) * inv;
        v = fmaf(v, g, b);
        out[idx] = v > 0.0f ? v : 0.0f;
    }
}

__global__ void bn_relu_eval_kernel(const float* __restrict__ x,
                                    float* __restrict__ out,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    const float* __restrict__ running_mean,
                                    const float* __restrict__ running_var,
                                    int total, int C, int L, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (idx / L) % C;
        float v = (x[idx] - running_mean[c]) * rsqrtf(running_var[c] + eps);
        v = fmaf(v, gamma[c], beta[c]);
        out[idx] = v > 0.0f ? v : 0.0f;
    }
}

__global__ void add_kernel(float* __restrict__ y, const float* __restrict__ x, int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        y[idx] += x[idx];
    }
}

torch::Tensor bn_relu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                           torch::Tensor running_mean, torch::Tensor running_var,
                           bool training, double momentum, double eps) {
    auto out = torch::empty_like(x);
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int L = (int)x.size(2);
    int total = N * C * L;
    const float* xp = x.data_ptr<float>();
    float* op = out.data_ptr<float>();
    const float* gp = weight.data_ptr<float>();
    const float* bp = bias.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        float* rmp = running_mean.defined() ? running_mean.data_ptr<float>() : nullptr;
        float* rvp = running_var.defined() ? running_var.data_ptr<float>() : nullptr;
        bn_relu_train_kernel<<<C, BN_THREADS, 0, stream>>>(xp, op, gp, bp, rmp, rvp, N, C, L,
                                                           (float)momentum, (float)eps);
    } else {
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 4096) blocks = 4096;
        bn_relu_eval_kernel<<<blocks, threads, 0, stream>>>(xp, op, gp, bp,
                                                            running_mean.data_ptr<float>(),
                                                            running_var.data_ptr<float>(),
                                                            total, C, L, (float)eps);
    }
    return out;
}

torch::Tensor add_residual_cuda(torch::Tensor y, torch::Tensor x) {
    int total = (int)y.numel();
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    add_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), x.data_ptr<float>(), total);
    return y;
}
"""

_ext = load_inline(
    name="dilated_prenorm_bn_relu_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bn_relu_forward", "add_residual"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, channels, kernel_size, dilation):
        super(ModelNew, self).__init__()
        self.norm = nn.BatchNorm1d(channels)
        self.activation = nn.ReLU()
        self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)

    def forward(self, x):
        z = _ext.bn_relu_forward(x, self.norm.weight, self.norm.bias,
                                 self.norm.running_mean, self.norm.running_var,
                                 bool(self.norm.training), float(self.norm.momentum), float(self.norm.eps))
        y = self.conv(z)
        return _ext.add_residual(y, x)