import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, divides by a constant, and applies LeakyReLU.
    """
    def __init__(self, in_channels, out_channels, kernel_size, divisor):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.divisor = divisor

    def forward(self, x):
        x = self.conv(x)
        x = x / self.divisor
        x = torch.nn.functional.leaky_relu(x, negative_slope=0.01)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
kernel_size = 3
divisor = 2

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, divisor]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor conv_div_lrelu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double divisor);
torch::Tensor conv_div_lrelu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double divisor) { return conv_div_lrelu_cuda(x, weight, bias, divisor); }
'''

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

__global__ void conv3x3_c8_o64_128_group16_kernel(const float* __restrict__ x,
                                               const float* __restrict__ w,
                                               const float* __restrict__ b,
                                               float* __restrict__ y,
                                               float inv_div) {
    constexpr int C = 8, H = 128, W = 128, OH = 126, OW = 126, HW = 16384;
    constexpr int OC = 64, SP = 15876;
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= SP) return;
    int g = blockIdx.y;
    int n = blockIdx.z;
    int oh = p / OW;
    int ow = p - oh * OW;
    int oc0 = g << 4;

    float s0=b[oc0+0], s1=b[oc0+1], s2=b[oc0+2], s3=b[oc0+3];
    float s4=b[oc0+4], s5=b[oc0+5], s6=b[oc0+6], s7=b[oc0+7];
    float s8=b[oc0+8], s9=b[oc0+9], s10=b[oc0+10], s11=b[oc0+11];
    float s12=b[oc0+12], s13=b[oc0+13], s14=b[oc0+14], s15=b[oc0+15];
    const int x_base_n = n * C * HW + oh * W + ow;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* xp = x + x_base_n + c * HW;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                float xv = xp[kh * W + kw];
                int wi = ((oc0 * C + c) * 3 + kh) * 3 + kw;
                s0=fmaf(xv,w[wi+0],s0);       s1=fmaf(xv,w[wi+72],s1);
                s2=fmaf(xv,w[wi+144],s2);     s3=fmaf(xv,w[wi+216],s3);
                s4=fmaf(xv,w[wi+288],s4);     s5=fmaf(xv,w[wi+360],s5);
                s6=fmaf(xv,w[wi+432],s6);     s7=fmaf(xv,w[wi+504],s7);
                s8=fmaf(xv,w[wi+576],s8);     s9=fmaf(xv,w[wi+648],s9);
                s10=fmaf(xv,w[wi+720],s10);   s11=fmaf(xv,w[wi+792],s11);
                s12=fmaf(xv,w[wi+864],s12);   s13=fmaf(xv,w[wi+936],s13);
                s14=fmaf(xv,w[wi+1008],s14);  s15=fmaf(xv,w[wi+1080],s15);
            }
        }
    }
#define ACT(v) do { v *= inv_div; v = v >= 0.0f ? v : v * 0.01f; } while(0)
    ACT(s0); ACT(s1); ACT(s2); ACT(s3); ACT(s4); ACT(s5); ACT(s6); ACT(s7);
    ACT(s8); ACT(s9); ACT(s10); ACT(s11); ACT(s12); ACT(s13); ACT(s14); ACT(s15);
#undef ACT
    int ybase = ((n * OC + oc0) * OH + oh) * OW + ow;
    y[ybase + 0 * SP]=s0;   y[ybase + 1 * SP]=s1;   y[ybase + 2 * SP]=s2;   y[ybase + 3 * SP]=s3;
    y[ybase + 4 * SP]=s4;   y[ybase + 5 * SP]=s5;   y[ybase + 6 * SP]=s6;   y[ybase + 7 * SP]=s7;
    y[ybase + 8 * SP]=s8;   y[ybase + 9 * SP]=s9;   y[ybase +10 * SP]=s10;  y[ybase +11 * SP]=s11;
    y[ybase +12 * SP]=s12;  y[ybase +13 * SP]=s13;  y[ybase +14 * SP]=s14;  y[ybase +15 * SP]=s15;
}

__global__ void conv3x3_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ b,
                                      float* __restrict__ y,
                                      int N, int C, int H, int W, int OC, int K,
                                      float inv_div) {
    int OH = H - K + 1, OW = W - K + 1;
    int64_t total = (int64_t)N * OC * OH * OW;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    int64_t q = idx / OW;
    int oh = q % OH; q /= OH;
    int oc = q % OC;
    int n = q / OC;
    float sum = b[oc];
    for (int c=0;c<C;++c) for(int kh=0;kh<K;++kh) for(int kw=0;kw<K;++kw) {
        sum = fmaf(x[((n*C+c)*H + oh+kh)*W + ow+kw], w[((oc*C+c)*K+kh)*K+kw], sum);
    }
    sum *= inv_div;
    y[idx] = sum >= 0.0f ? sum : sum * 0.01f;
}

torch::Tensor conv_div_lrelu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double divisor) {
    int N=(int)x.size(0), C=(int)x.size(1), H=(int)x.size(2), W=(int)x.size(3);
    int OC=(int)weight.size(0), K=(int)weight.size(2);
    int OH=H-K+1, OW=W-K+1;
    auto y = torch::empty({N, OC, OH, OW}, x.options());
    float inv = 1.0f / (float)divisor;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int threads = 256;
    if (C == 8 && OC == 64 && K == 3 && H == 128 && W == 128) {
        dim3 grid((15876 + threads - 1) / threads, 4, N);
        conv3x3_c8_o64_128_group16_kernel<<<grid, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), inv);
    } else {
        int64_t total = (int64_t)N * OC * OH * OW;
        int blocks = (int)((total + threads - 1) / threads);
        conv3x3_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, OC, K, inv);
    }
    return y;
}
'''

_ext = load_inline(name='conv_div_lrelu_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['conv_div_lrelu_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, divisor):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.divisor = divisor
    def forward(self, x):
        return _ext.conv_div_lrelu_forward(x, self.conv.weight, self.conv.bias, float(self.divisor))