import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Pre-norm 1D dilated convolution block: LayerNorm, GELU activation, dilated Conv1d, and residual addition.
    Common in genomic sequence models (Enformer, Basenji) for capturing long-range dependencies in DNA/RNA sequences.
    """
    def __init__(self, channels, seq_len, kernel_size, dilation):
        super(Model, self).__init__()
        self.norm = nn.LayerNorm([channels, seq_len])
        self.activation = nn.GELU()
        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 = 4

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

def get_init_inputs():
    return [channels, seq_len, 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 ln_gelu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps);
torch::Tensor add_forward(torch::Tensor y, torch::Tensor residual);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define THREADS 256
#define CHUNKS 256

__inline__ __device__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__inline__ __device__ float block_reduce_sum(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) v = warp_reduce_sum(v);
    return v;
}

__global__ void partial_reduce_kernel(const float* __restrict__ x,
                                      float* __restrict__ partial,
                                      int B, int N) {
    int chunk = blockIdx.x;
    int b = blockIdx.y;
    int chunk_size = (N + CHUNKS - 1) / CHUNKS;
    int start = chunk * chunk_size;
    int end = start + chunk_size;
    if (end > N) end = N;

    float s = 0.0f;
    float ss = 0.0f;
    const float* xb = x + ((long long)b) * N;
    for (int i = start + threadIdx.x; i < end; i += blockDim.x) {
        float v = xb[i];
        s += v;
        ss = fmaf(v, v, ss);
    }
    s = block_reduce_sum(s);
    ss = block_reduce_sum(ss);
    if (threadIdx.x == 0) {
        int idx = (b * CHUNKS + chunk) * 2;
        partial[idx] = s;
        partial[idx + 1] = ss;
    }
}

__global__ void finalize_kernel(const float* __restrict__ partial,
                                float* __restrict__ stats,
                                int B, int N, float eps) {
    int b = blockIdx.x;
    float s = 0.0f;
    float ss = 0.0f;
    for (int c = threadIdx.x; c < CHUNKS; c += blockDim.x) {
        int idx = (b * CHUNKS + c) * 2;
        s += partial[idx];
        ss += partial[idx + 1];
    }
    s = block_reduce_sum(s);
    ss = block_reduce_sum(ss);
    if (threadIdx.x == 0) {
        float invN = 1.0f / (float)N;
        float mean = s * invN;
        float var = ss * invN - mean * mean;
        if (var < 0.0f) var = 0.0f;
        stats[b * 2] = mean;
        stats[b * 2 + 1] = rsqrtf(var + eps);
    }
}

__global__ void apply_ln_gelu_kernel(const float* __restrict__ x,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     const float* __restrict__ stats,
                                     float* __restrict__ out,
                                     long long total, int N) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float inv_sqrt2 = 0.70710678118654752440f;
    for (; idx < total; idx += stride) {
        int b = (int)(idx / N);
        int n = (int)(idx - (long long)b * N);
        float z = (x[idx] - stats[b * 2]) * stats[b * 2 + 1];
        z = z * gamma[n] + beta[n];
        float g = 0.5f * z * (1.0f + erff(z * inv_sqrt2));
        out[idx] = g;
    }
}

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

torch::Tensor ln_gelu_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
    x = x.contiguous();
    gamma = gamma.contiguous();
    beta = beta.contiguous();
    auto out = torch::empty_like(x);
    int B = (int)x.size(0);
    int N = (int)(x.numel() / B);
    auto partial = torch::empty({B, CHUNKS, 2}, x.options());
    auto stats = torch::empty({B, 2}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 grid1(CHUNKS, B);
    partial_reduce_kernel<<<grid1, THREADS, 0, stream>>>(x.data_ptr<float>(), partial.data_ptr<float>(), B, N);
    finalize_kernel<<<B, THREADS, 0, stream>>>(partial.data_ptr<float>(), stats.data_ptr<float>(), B, N, (float)eps);
    long long total = x.numel();
    int blocks = (int)((total + THREADS - 1) / THREADS);
    if (blocks > 65535) blocks = 65535;
    apply_ln_gelu_kernel<<<blocks, THREADS, 0, stream>>>(x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), stats.data_ptr<float>(), out.data_ptr<float>(), total, N);
    return out;
}

torch::Tensor add_forward(torch::Tensor y, torch::Tensor residual) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(y));
    y = y.contiguous();
    residual = residual.contiguous();
    auto out = torch::empty_like(y);
    long long total = y.numel();
    int blocks = (int)((total + THREADS - 1) / THREADS);
    if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    add_kernel<<<blocks, THREADS, 0, stream>>>(y.data_ptr<float>(), residual.data_ptr<float>(), out.data_ptr<float>(), total);
    return out;
}
'''

_ext = load_inline(
    name='prenorm_gelu_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['ln_gelu_forward', 'add_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, channels, seq_len, kernel_size, dilation):
        super(ModelNew, self).__init__()
        self.norm = nn.LayerNorm([channels, seq_len])
        self.activation = nn.GELU()
        self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)

    def forward(self, x):
        residual = x
        x = _ext.ln_gelu_forward(x, self.norm.weight, self.norm.bias, float(self.norm.eps))
        x = self.conv(x)
        x = _ext.add_forward(x, residual)
        return x