import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        """
        Initialize the VGG16 model.
        
        :param num_classes: The number of output classes (default is 1000 for ImageNet)
        """
        super(Model, self).__init__()
        
        # VGG16 architecture: 5 blocks of convolutional layers followed by max pooling
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            # Block 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            # Block 4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            # Block 5
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # Fully connected layers
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.0),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.0),
            nn.Linear(4096, num_classes)
        )
    
    def forward(self, x):
        """
        Forward pass of the VGG16 model.
        
        :param x: The input tensor, shape (batch_size, 3, 224, 224)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# Test code
batch_size = 10
num_classes = 1000

def get_inputs():
    return [torch.rand(batch_size, 3, 224, 224)]

def get_init_inputs():
    return [num_classes]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor relu_maxpool2x2_cl_cuda(torch::Tensor input);
'''

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

__global__ void relu_maxpool2x2_cl_kernel(const float* __restrict__ inp,
                                          float* __restrict__ out,
                                          int total, int C, int H, int W) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int OW = W >> 1;
    int OH = H >> 1;

    // Output tensor is channels-last physically: (((n*OH + oy)*OW + ox)*C + c)
    int c = idx % C;
    int t = idx / C;
    int ox = t % OW;
    t /= OW;
    int oy = t % OH;
    int n = t / OH;

    int in_base = ((n * H + (oy << 1)) * W + (ox << 1)) * C + c;
    float a = inp[in_base];
    float b = inp[in_base + C];
    float cc = inp[in_base + (long long)W * C];
    float d = inp[in_base + ((long long)W + 1) * C];
    float m = fmaxf(fmaxf(a, b), fmaxf(cc, d));
    out[idx] = fmaxf(m, 0.0f);
}

torch::Tensor relu_maxpool2x2_cl_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");

    auto N = input.size(0);
    auto C = input.size(1);
    auto H = input.size(2);
    auto W = input.size(3);
    auto output = torch::empty_strided({N, C, H/2, W/2},
                                      {C*(H/2)*(W/2), 1, (W/2)*C, C},
                                      input.options());

    int64_t total64 = N * C * (H/2) * (W/2);
    int threads = 256;
    int blocks = (int)((total64 + threads - 1) / threads);
    relu_maxpool2x2_cl_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), (int)total64, (int)C, (int)H, (int)W);
    return output;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.0),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.0),
            nn.Linear(4096, num_classes),
        )
        # Keep convolution operands in NHWC physical layout for faster cuDNN kernels.
        self.features.to(memory_format=torch.channels_last)

    def forward(self, x):
        f = self.features
        x = x.contiguous(memory_format=torch.channels_last)

        x = torch.relu_(f[0](x))
        x = f[2](x)
        x = _ext.relu_maxpool2x2_cl_cuda(x)

        x = torch.relu_(f[5](x))
        x = f[7](x)
        x = _ext.relu_maxpool2x2_cl_cuda(x)

        x = torch.relu_(f[10](x))
        x = torch.relu_(f[12](x))
        x = f[14](x)
        x = _ext.relu_maxpool2x2_cl_cuda(x)

        x = torch.relu_(f[17](x))
        x = torch.relu_(f[19](x))
        x = f[21](x)
        x = _ext.relu_maxpool2x2_cl_cuda(x)

        x = torch.relu_(f[24](x))
        x = torch.relu_(f[26](x))
        x = f[28](x)
        x = _ext.relu_maxpool2x2_cl_cuda(x)

        x = torch.flatten(x, 1)
        c = self.classifier
        x = torch.relu_(c[0](x))
        x = torch.relu_(c[3](x))
        x = c[6](x)
        return x