import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        """
        Initialize the VGG19 model.

        :param num_classes: The number of output classes (default is 1000 for ImageNet)
        """
        super(Model, self).__init__()
        
        # VGG19 architecture: 16 Conv layers + 5 MaxPool layers + 3 Fully Connected layers
        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.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.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.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)
        )
    
    def forward(self, x):
        """
        Forward pass of the VGG19 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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor flatten_nchw(torch::Tensor x);
"""

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

__global__ void flatten_nchw_kernel(const float* __restrict__ x, float* __restrict__ y,
                                    int N, int C, int H, int W) {
    int64_t total = (int64_t)N * C * H * W;
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int w = idx % W;
        int t = idx / W;
        int h = t % H;
        t /= H;
        int c = t % C;
        int n = t / C;
        // input is channels_last contiguous (N,H,W,C), output is NCHW-flattened
        y[idx] = x[(int64_t)n * H * W * C + (int64_t)h * W * C + (int64_t)w * C + c];
    }
}

torch::Tensor flatten_nchw(torch::Tensor x) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    auto y = torch::empty({N, C * H * W}, x.options());
    int64_t total = (int64_t)N * C * H * W;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 1024) blocks = 1024;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    flatten_nchw_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}
"""

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

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.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.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.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)
        )
        self.features.to(memory_format=torch.channels_last)
        self._graph = None
        self._static_x = None
        self._static_y = None
        self._graph_shape = None

    def _impl(self, x):
        x = self.features(x)
        x = _ext.flatten_nchw(x)
        x = F.linear(x, self.classifier[0].weight, self.classifier[0].bias)
        x = F.relu(x, inplace=True)
        x = F.linear(x, self.classifier[3].weight, self.classifier[3].bias)
        x = F.relu(x, inplace=True)
        return F.linear(x, self.classifier[6].weight, self.classifier[6].bias)

    def _capture(self, x):
        self._static_x = torch.empty_like(x, memory_format=torch.channels_last)
        self._static_x.copy_(x)
        with torch.no_grad():
            _ = self._impl(self._static_x)  # warm up cuDNN/cuBLAS and extension allocation paths
            torch.cuda.synchronize()
            self._graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(self._graph):
                self._static_y = self._impl(self._static_x)
            self._graph.replay()
        self._graph_shape = tuple(x.shape)
        return self._static_y.clone()

    def forward(self, x):
        if x.is_cuda and x.dtype is torch.float32:
            if self._graph is None or self._graph_shape != tuple(x.shape):
                return self._capture(x)
            self._static_x.copy_(x)
            self._graph.replay()
            return self._static_y.clone()
        with torch.no_grad():
            return self._impl(x.contiguous(memory_format=torch.channels_last))