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

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        """
        :param num_classes: The number of output classes (default is 1000 for ImageNet)
        """
        super(Model, self).__init__()
        
        # First convolutional layer
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2)
        self.relu1 = nn.ReLU(inplace=True)
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)
        
        # Second convolutional layer
        self.conv2 = nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2)
        self.relu2 = nn.ReLU(inplace=True)
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2)
        
        # Third convolutional layer
        self.conv3 = nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1)
        self.relu3 = nn.ReLU(inplace=True)
        
        # Fourth convolutional layer
        self.conv4 = nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1)
        self.relu4 = nn.ReLU(inplace=True)
        
        # Fifth convolutional layer
        self.conv5 = nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1)
        self.relu5 = nn.ReLU(inplace=True)
        self.maxpool3 = nn.MaxPool2d(kernel_size=3, stride=2)
        
        # Fully connected layers
        self.fc1 = nn.Linear(in_features=256 * 6 * 6, out_features=4096)
        self.relu6 = nn.ReLU(inplace=True)
        self.dropout1 = nn.Dropout(p=0.0)
        
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.relu7 = nn.ReLU(inplace=True)
        self.dropout2 = nn.Dropout(p=0.0)
        
        self.fc3 = nn.Linear(in_features=4096, out_features=num_classes)
    
    def forward(self, x):
        """
        :param x: The input tensor, shape (batch_size, 3, 224, 224)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.maxpool1(x)
        
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.maxpool2(x)
        
        x = self.conv3(x)
        x = self.relu3(x)
        
        x = self.conv4(x)
        x = self.relu4(x)
        
        x = self.conv5(x)
        x = self.relu5(x)
        x = self.maxpool3(x)
        
        x = torch.flatten(x, 1)
        
        x = self.fc1(x)
        x = self.relu6(x)
        x = self.dropout1(x)
        
        x = self.fc2(x)
        x = self.relu7(x)
        x = self.dropout2(x)
        
        x = self.fc3(x)
        
        return x

# Test code
batch_size = 1024
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 relu_pool3x3s2_forward(torch::Tensor x);
torch::Tensor relu_inplace_forward(torch::Tensor x);
"""

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

__global__ void relu_pool3x3s2_kernel(const float* __restrict__ x,
                                      float* __restrict__ y,
                                      int total, int C, int H, int W,
                                      int outH, int outW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = idx % outW;
    int t = idx / outW;
    int oh = t % outH;
    t /= outH;
    int c = t % C;
    int n = t / C;

    int ih = oh * 2;
    int iw = ow * 2;
    int base = ((n * C + c) * H + ih) * W + iw;

    float m = x[base];
    float v;
    v = x[base + 1];       m = v > m ? v : m;
    v = x[base + 2];       m = v > m ? v : m;
    int b1 = base + W;
    v = x[b1];             m = v > m ? v : m;
    v = x[b1 + 1];         m = v > m ? v : m;
    v = x[b1 + 2];         m = v > m ? v : m;
    int b2 = base + 2 * W;
    v = x[b2];             m = v > m ? v : m;
    v = x[b2 + 1];         m = v > m ? v : m;
    v = x[b2 + 2];         m = v > m ? v : m;

    y[idx] = m > 0.0f ? m : 0.0f;
}

__global__ void relu_inplace_kernel(float* __restrict__ x, int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        float v = x[idx];
        x[idx] = v > 0.0f ? v : 0.0f;
    }
}

torch::Tensor relu_pool3x3s2_forward(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    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 outH = (H - 3) / 2 + 1;
    const int outW = (W - 3) / 2 + 1;
    auto y = torch::empty({N, C, outH, outW}, x.options());
    const int total = N * C * outH * outW;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    relu_pool3x3s2_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), total, C, H, W, outH, outW);
    return y;
}

torch::Tensor relu_inplace_forward(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    const int total = (int)x.numel();
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    relu_inplace_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), total);
    return x;
}
"""

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


class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2)
        self.relu1 = nn.ReLU(inplace=True)
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv2 = nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2)
        self.relu2 = nn.ReLU(inplace=True)
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv3 = nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1)
        self.relu3 = nn.ReLU(inplace=True)
        self.conv4 = nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1)
        self.relu4 = nn.ReLU(inplace=True)
        self.conv5 = nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1)
        self.relu5 = nn.ReLU(inplace=True)
        self.maxpool3 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.fc1 = nn.Linear(in_features=256 * 6 * 6, out_features=4096)
        self.relu6 = nn.ReLU(inplace=True)
        self.dropout1 = nn.Dropout(p=0.0)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.relu7 = nn.ReLU(inplace=True)
        self.dropout2 = nn.Dropout(p=0.0)
        self.fc3 = nn.Linear(in_features=4096, out_features=num_classes)

    def forward(self, x):
        x = F.conv2d(x, self.conv1.weight, self.conv1.bias, stride=4, padding=2)
        x = _ext.relu_pool3x3s2_forward(x)

        x = F.conv2d(x, self.conv2.weight, self.conv2.bias, stride=1, padding=2)
        x = _ext.relu_pool3x3s2_forward(x)

        x = F.conv2d(x, self.conv3.weight, self.conv3.bias, stride=1, padding=1)
        x = _ext.relu_inplace_forward(x)

        x = F.conv2d(x, self.conv4.weight, self.conv4.bias, stride=1, padding=1)
        x = _ext.relu_inplace_forward(x)

        x = F.conv2d(x, self.conv5.weight, self.conv5.bias, stride=1, padding=1)
        x = _ext.relu_pool3x3s2_forward(x)

        x = torch.flatten(x, 1)
        x = F.linear(x, self.fc1.weight, self.fc1.bias)
        x = _ext.relu_inplace_forward(x)
        x = F.linear(x, self.fc2.weight, self.fc2.bias)
        x = _ext.relu_inplace_forward(x)
        x = F.linear(x, self.fc3.weight, self.fc3.bias)
        return x