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

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

__device__ __forceinline__ void warpReduceSum(volatile float* sdata, int tid) {
    sdata[tid] += sdata[tid + 32];
    sdata[tid] += sdata[tid + 16];
    sdata[tid] += sdata[tid + 8];
    sdata[tid] += sdata[tid + 4];
    sdata[tid] += sdata[tid + 2];
    sdata[tid] += sdata[tid + 1];
}

__global__ void ln_gelu_kernel(const float* __restrict__ x, float* __restrict__ y,
                               const float* __restrict__ gamma, const float* __restrict__ beta,
                               int B, int N, float eps) {
    int b = blockIdx.x;
    __shared__ float sdata[256];
    int tid = threadIdx.x;
    double sum = 0.0;
    for (int i = tid; i < N; i += 256) sum += (double)x[b * N + i];
    sdata[tid] = (float)sum;
    __syncthreads();
    for (int s = 128; s > 32; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid < 32) warpReduceSum(sdata, tid);
    __syncthreads();
    float mean = sdata[0] / N;
    double varsum = 0.0;
    for (int i = tid; i < N; i += 256) {
        double v = (double)x[b*N + i] - mean;
        varsum += v * v;
    }
    sdata[tid] = (float)varsum;
    __syncthreads();
    for (int s = 128; s > 32; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid < 32) warpReduceSum(sdata, tid);
    __syncthreads();
    float var = sdata[0] / N;
    float invstd = rsqrtf(var + eps);
    for (int i = tid; i < N; i += 256) {
        float v = (x[b*N + i] - mean) * invstd;
        v = v * gamma[i] + beta[i];
        float cdf = 0.5f * (1.0f + erff(v * 0.70710678118f));
        v = v * cdf;
        y[b*N + i] = v;
    }
}

torch::Tensor fused_ln_gelu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta) {
    int B = x.size(0);
    int N = x.size(1) * x.size(2);
    auto y = torch::empty_like(x);
    ln_gelu_kernel<<<B, 256>>>(x.data_ptr<float>(), y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), B, N, 1e-5f);
    return y;
}
'''

cpp_source = '''
torch::Tensor fused_ln_gelu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta);
'''

module = load_inline(name='fused_mod3', cpp_sources=[cpp_source], cuda_sources=[cuda_source], functions=['fused_ln_gelu'], 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.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)
        self.m = module

    def forward(self, x):
        residual = x
        x = self.m.fused_ln_gelu(x, self.norm.weight, self.norm.bias)
        x = self.conv(x)
        return x + residual