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

torch.backends.cudnn.benchmark = True

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor relu_maxpool3x3s2_cl(torch::Tensor x);
torch::Tensor relu_maxpool3x3s2_cl_to_nchw(torch::Tensor x);
torch::Tensor relu_inplace(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_kernel(float* __restrict__ x, int64_t n) {
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < n; idx += stride) {
        float v = x[idx];
        x[idx] = v > 0.f ? v : 0.f;
    }
}

__global__ void relu_maxpool3x3s2_cl_kernel(const float* __restrict__ x, float* __restrict__ y,
                                            int total, int C, int H, int W, int Ho, int Wo) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = idx % C;
        int t = idx / C;
        int ow = t % Wo;
        t /= Wo;
        int oh = t % Ho;
        int n = t / Ho;
        int base = ((n * H + oh * 2) * W + ow * 2) * C + c;
        float m = 0.0f;
        float v;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        base += W*C;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        base += W*C;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        y[idx] = m > 0.f ? m : 0.f;
    }
}

__global__ void relu_maxpool3x3s2_cl_to_nchw_kernel(const float* __restrict__ x, float* __restrict__ y,
                                                    int total, int C, int H, int W, int Ho, int Wo) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int ow = idx % Wo;
        int t = idx / Wo;
        int oh = t % Ho;
        t /= Ho;
        int c = t % C;
        int n = t / C;
        int base = ((n * H + oh * 2) * W + ow * 2) * C + c;
        float m = 0.0f;
        float v;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        base += W*C;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        base += W*C;
        v = x[base]; m = fmaxf(m, v);
        v = x[base + C]; m = fmaxf(m, v);
        v = x[base + 2*C]; m = fmaxf(m, v);
        y[idx] = m > 0.f ? m : 0.f;
    }
}

torch::Tensor relu_maxpool3x3s2_cl(torch::Tensor x) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int Ho = (H - 3) / 2 + 1, Wo = (W - 3) / 2 + 1;
    auto y = torch::empty_strided({N, C, Ho, Wo}, {C*Ho*Wo, 1, C*Wo, C}, x.options());
    int total = N*C*Ho*Wo;
    const int threads = 256;
    int blocks = (total + threads - 1) / threads; if (blocks > 65535) blocks = 65535;
    relu_maxpool3x3s2_cl_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), total, C, H, W, Ho, Wo);
    return y;
}

torch::Tensor relu_maxpool3x3s2_cl_to_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);
    int Ho = (H - 3) / 2 + 1, Wo = (W - 3) / 2 + 1;
    auto y = torch::empty({N, C, Ho, Wo}, x.options());
    int total = N*C*Ho*Wo;
    const int threads = 256;
    int blocks = (total + threads - 1) / threads; if (blocks > 65535) blocks = 65535;
    relu_maxpool3x3s2_cl_to_nchw_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), total, C, H, W, Ho, Wo);
    return y;
}

torch::Tensor relu_inplace(torch::Tensor x) {
    int64_t n = x.numel();
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads); if (blocks > 65535) blocks = 65535;
    relu_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), n);
    return x;
}
'''
_ext = load_inline(name='alexnet_fused_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['relu_maxpool3x3s2_cl', 'relu_maxpool3x3s2_cl_to_nchw', 'relu_inplace'],
                   verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2)
        self.conv2 = nn.Conv2d(96, 256, kernel_size=5, padding=2)
        self.conv3 = nn.Conv2d(256, 384, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(384, 384, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(384, 256, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(256 * 6 * 6, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)
        for m in (self.conv1, self.conv2, self.conv3, self.conv4, self.conv5):
            m.weight.data = m.weight.data.contiguous(memory_format=torch.channels_last)
    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last)
        x = F.conv2d(x, self.conv1.weight, self.conv1.bias, stride=4, padding=2)
        x = _ext.relu_maxpool3x3s2_cl(x)
        x = F.conv2d(x, self.conv2.weight, self.conv2.bias, padding=2)
        x = _ext.relu_maxpool3x3s2_cl(x)
        x = F.conv2d(x, self.conv3.weight, self.conv3.bias, padding=1)
        x = _ext.relu_inplace(x)
        x = F.conv2d(x, self.conv4.weight, self.conv4.bias, padding=1)
        x = _ext.relu_inplace(x)
        x = F.conv2d(x, self.conv5.weight, self.conv5.bias, padding=1)
        x = _ext.relu_maxpool3x3s2_cl_to_nchw(x)
        x = torch.flatten(x, 1)
        x = F.linear(x, self.fc1.weight, self.fc1.bias)
        x = _ext.relu_inplace(x)
        x = F.linear(x, self.fc2.weight, self.fc2.bias)
        x = _ext.relu_inplace(x)
        x = F.linear(x, self.fc3.weight, self.fc3.bias)
        return x