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 fused_conv_div_lrelu_forward(torch::Tensor x,
                                           torch::Tensor weight,
                                           torch::Tensor bias,
                                           double divisor);
"""

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

#define TILE 16
#define IN_TILE 18
#define CIN_FIXED 8
#define K_FIXED 3
#define OCG 8

__global__ __launch_bounds__(256, 2)
void conv3x3_c8_ocg8_kernel(const float* __restrict__ x,
                            const float* __restrict__ w,
                            const float* __restrict__ b,
                            float* __restrict__ out,
                            int N, int Cout, int H, int W,
                            int OH, int OW,
                            float inv_divisor) {
    __shared__ float sh[CIN_FIXED * IN_TILE * IN_TILE];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * TILE + tx;

    const int tile_ow = blockIdx.x * TILE;
    const int tile_oh = blockIdx.y * TILE;
    const int oc_groups = (Cout + OCG - 1) / OCG;
    const int group = blockIdx.z % oc_groups;
    const int n = blockIdx.z / oc_groups;
    const int oc0 = group * OCG;

    // Cooperative load of the input patch needed by this 16x16 output tile.
    #pragma unroll
    for (int idx = tid; idx < CIN_FIXED * IN_TILE * IN_TILE; idx += TILE * TILE) {
        int rem = idx;
        const int ci = rem / (IN_TILE * IN_TILE);
        rem -= ci * (IN_TILE * IN_TILE);
        const int sy = rem / IN_TILE;
        const int sx = rem - sy * IN_TILE;
        const int ih = tile_oh + sy;
        const int iw = tile_ow + sx;
        float v = 0.0f;
        if (ih < H && iw < W) {
            v = x[((n * CIN_FIXED + ci) * H + ih) * W + iw];
        }
        sh[idx] = v;
    }
    __syncthreads();

    const int oh = tile_oh + ty;
    const int ow = tile_ow + tx;
    if (oh >= OH || ow >= OW) return;

    float acc0 = (oc0 + 0 < Cout) ? b[oc0 + 0] : 0.0f;
    float acc1 = (oc0 + 1 < Cout) ? b[oc0 + 1] : 0.0f;
    float acc2 = (oc0 + 2 < Cout) ? b[oc0 + 2] : 0.0f;
    float acc3 = (oc0 + 3 < Cout) ? b[oc0 + 3] : 0.0f;
    float acc4 = (oc0 + 4 < Cout) ? b[oc0 + 4] : 0.0f;
    float acc5 = (oc0 + 5 < Cout) ? b[oc0 + 5] : 0.0f;
    float acc6 = (oc0 + 6 < Cout) ? b[oc0 + 6] : 0.0f;
    float acc7 = (oc0 + 7 < Cout) ? b[oc0 + 7] : 0.0f;

    #pragma unroll
    for (int ci = 0; ci < CIN_FIXED; ++ci) {
        #pragma unroll
        for (int ky = 0; ky < K_FIXED; ++ky) {
            #pragma unroll
            for (int kx = 0; kx < K_FIXED; ++kx) {
                const float xv = sh[(ci * IN_TILE + (ty + ky)) * IN_TILE + (tx + kx)];
                const int kbase = ((ci * K_FIXED + ky) * K_FIXED + kx);
                if (oc0 + 0 < Cout) acc0 = fmaf(xv, w[((oc0 + 0) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc0);
                if (oc0 + 1 < Cout) acc1 = fmaf(xv, w[((oc0 + 1) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc1);
                if (oc0 + 2 < Cout) acc2 = fmaf(xv, w[((oc0 + 2) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc2);
                if (oc0 + 3 < Cout) acc3 = fmaf(xv, w[((oc0 + 3) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc3);
                if (oc0 + 4 < Cout) acc4 = fmaf(xv, w[((oc0 + 4) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc4);
                if (oc0 + 5 < Cout) acc5 = fmaf(xv, w[((oc0 + 5) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc5);
                if (oc0 + 6 < Cout) acc6 = fmaf(xv, w[((oc0 + 6) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc6);
                if (oc0 + 7 < Cout) acc7 = fmaf(xv, w[((oc0 + 7) * CIN_FIXED * K_FIXED * K_FIXED) + kbase], acc7);
            }
        }
    }

    const int spatial = OH * OW;
    const int out_base = (n * Cout * spatial) + oh * OW + ow;

    if (oc0 + 0 < Cout) { float v = acc0 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 0) * spatial] = v; }
    if (oc0 + 1 < Cout) { float v = acc1 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 1) * spatial] = v; }
    if (oc0 + 2 < Cout) { float v = acc2 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 2) * spatial] = v; }
    if (oc0 + 3 < Cout) { float v = acc3 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 3) * spatial] = v; }
    if (oc0 + 4 < Cout) { float v = acc4 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 4) * spatial] = v; }
    if (oc0 + 5 < Cout) { float v = acc5 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 5) * spatial] = v; }
    if (oc0 + 6 < Cout) { float v = acc6 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 6) * spatial] = v; }
    if (oc0 + 7 < Cout) { float v = acc7 * inv_divisor; v = (v > 0.0f) ? v : v * 0.01f; out[out_base + (oc0 + 7) * spatial] = v; }
}

torch::Tensor fused_conv_div_lrelu_forward(torch::Tensor x,
                                           torch::Tensor weight,
                                           torch::Tensor bias,
                                           double divisor) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int Cout = (int)weight.size(0);
    const int K = (int)weight.size(2);
    TORCH_CHECK(C == CIN_FIXED && K == K_FIXED && weight.size(3) == K_FIXED,
                "this optimized kernel supports only C=8, K=3");

    const int OH = H - K_FIXED + 1;
    const int OW = W - K_FIXED + 1;
    auto out = torch::empty({N, Cout, OH, OW}, x.options());

    dim3 block(TILE, TILE, 1);
    dim3 grid((OW + TILE - 1) / TILE,
              (OH + TILE - 1) / TILE,
              N * ((Cout + OCG - 1) / OCG));

    const float inv_div = 1.0f / (float)divisor;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3x3_c8_ocg8_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        N, Cout, H, W, OH, OW, inv_div);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="fused_conv_div_lrelu_c8k3_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_conv_div_lrelu_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """Fused Conv2d / divisor / LeakyReLU implementation for the given model."""
    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.fused_conv_div_lrelu_forward(x, self.conv.weight, self.conv.bias, float(self.divisor))