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

class Model(nn.Module):
    def __init__(self, num_classes):
        """
        LeNet-5 architecture implementation in PyTorch.

        :param num_classes: The number of output classes.
        """
        super(Model, self).__init__()
        
        # Convolutional layers
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        
        # Fully connected layers
        self.fc1 = nn.Linear(in_features=16*5*5, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=num_classes)
    
    def forward(self, x):
        """
        Forward pass of the LeNet-5 model.

        :param x: The input tensor, shape (batch_size, 1, 32, 32)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        # First convolutional layer with ReLU activation and max pooling
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        
        # Second convolutional layer with ReLU activation and max pooling
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        
        # Flatten the output for the fully connected layers
        x = x.view(-1, 16*5*5)
        
        # First fully connected layer with ReLU activation
        x = F.relu(self.fc1(x))
        
        # Second fully connected layer with ReLU activation
        x = F.relu(self.fc2(x))
        
        # Final fully connected layer
        x = self.fc3(x)
        
        return x

# Test code for the LeNet-5 model (larger batch & image)
batch_size = 4096
num_classes = 20

def get_inputs():
    return [torch.rand(batch_size, 1, 32, 32)]

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 lenet_forward(torch::Tensor x, torch::Tensor c1w, torch::Tensor c1b, torch::Tensor c2w, torch::Tensor c2b, torch::Tensor f1w, torch::Tensor f1b, torch::Tensor f2w, torch::Tensor f2b, torch::Tensor f3w, torch::Tensor f3b);
"""

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

__global__ void conv1_pool_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ y, int B) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * 6 * 14 * 14;
    if (idx >= total) return;
    int t = idx; int px = t % 14; t /= 14; int py = t % 14; t /= 14; int oc = t % 6; t /= 6; int n = t;
    float m = -FLT_MAX; int base_x = n * 1024; int base_w = oc * 25;
    #pragma unroll
    for (int dy=0; dy<2; ++dy) {
        #pragma unroll
        for (int dx=0; dx<2; ++dx) {
            int oy = py * 2 + dy; int ox = px * 2 + dx; float s = b[oc];
            #pragma unroll
            for (int ky=0; ky<5; ++ky) {
                #pragma unroll
                for (int kx=0; kx<5; ++kx) s = fmaf(x[base_x + (oy+ky)*32 + (ox+kx)], w[base_w + ky*5 + kx], s);
            }
            m = fmaxf(m, s);
        }
    }
    y[idx] = fmaxf(m, 0.0f);
}

__global__ void conv2_pool_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ y, int B) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * 16 * 5 * 5;
    if (idx >= total) return;
    int t = idx; int px = t % 5; t /= 5; int py = t % 5; t /= 5; int oc = t % 16; t /= 16; int n = t;
    float m = -FLT_MAX; int base_in_n = n * 1176; int base_w_oc = oc * 150;
    #pragma unroll
    for (int dy=0; dy<2; ++dy) {
        #pragma unroll
        for (int dx=0; dx<2; ++dx) {
            int oy = py * 2 + dy; int ox = px * 2 + dx; float s = b[oc];
            #pragma unroll
            for (int ic=0; ic<6; ++ic) {
                int base_in = base_in_n + ic * 196; int base_w = base_w_oc + ic * 25;
                #pragma unroll
                for (int ky=0; ky<5; ++ky) {
                    #pragma unroll
                    for (int kx=0; kx<5; ++kx) s = fmaf(x[base_in + (oy+ky)*14 + (ox+kx)], w[base_w + ky*5 + kx], s);
                }
            }
            m = fmaxf(m, s);
        }
    }
    y[idx] = fmaxf(m, 0.0f);
}

#define TM 16
#define TN 16
#define TK 32
__global__ void linear_tiled_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ y, int B, int K, int N, int relu) {
    __shared__ float As[TM][TK];
    __shared__ float Bs[TK][TN];
    int row = blockIdx.y * TM + threadIdx.y;
    int col = blockIdx.x * TN + threadIdx.x;
    float acc = (col < N ? b[col] : 0.0f);
    int tid = threadIdx.y * TN + threadIdx.x;
    for (int kk = 0; kk < K; kk += TK) {
        for (int i = tid; i < TM*TK; i += TM*TN) {
            int r = i / TK; int k = i - r * TK; int gr = blockIdx.y * TM + r; int gk = kk + k;
            As[r][k] = (gr < B && gk < K) ? x[gr*K + gk] : 0.0f;
        }
        for (int i = tid; i < TK*TN; i += TM*TN) {
            int k = i / TN; int c = i - k * TN; int gc = blockIdx.x * TN + c; int gk = kk + k;
            Bs[k][c] = (gc < N && gk < K) ? w[gc*K + gk] : 0.0f;
        }
        __syncthreads();
        #pragma unroll
        for (int k = 0; k < TK; ++k) acc = fmaf(As[threadIdx.y][k], Bs[k][threadIdx.x], acc);
        __syncthreads();
    }
    if (row < B && col < N) { if (relu) acc = fmaxf(acc, 0.0f); y[row*N + col] = acc; }
}

torch::Tensor lenet_forward(torch::Tensor x, torch::Tensor c1w, torch::Tensor c1b, torch::Tensor c2w, torch::Tensor c2b, torch::Tensor f1w, torch::Tensor f1b, torch::Tensor f2w, torch::Tensor f2b, torch::Tensor f3w, torch::Tensor f3b) {
    int B = x.size(0); int C = f3w.size(0); auto opts = x.options();
    auto y1 = torch::empty({B,6,14,14}, opts); auto y2 = torch::empty({B,16,5,5}, opts); auto y3 = torch::empty({B,120}, opts); auto y4 = torch::empty({B,84}, opts); auto out = torch::empty({B,C}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream(); int threads = 256;
    conv1_pool_kernel<<<(B*6*14*14 + threads - 1)/threads, threads, 0, stream>>>(x.data_ptr<float>(), c1w.data_ptr<float>(), c1b.data_ptr<float>(), y1.data_ptr<float>(), B);
    conv2_pool_kernel<<<(B*16*5*5 + threads - 1)/threads, threads, 0, stream>>>(y1.data_ptr<float>(), c2w.data_ptr<float>(), c2b.data_ptr<float>(), y2.data_ptr<float>(), B);
    dim3 block(TN, TM);
    linear_tiled_kernel<<<dim3((120+TN-1)/TN, (B+TM-1)/TM), block, 0, stream>>>(y2.data_ptr<float>(), f1w.data_ptr<float>(), f1b.data_ptr<float>(), y3.data_ptr<float>(), B, 400, 120, 1);
    linear_tiled_kernel<<<dim3((84+TN-1)/TN, (B+TM-1)/TM), block, 0, stream>>>(y3.data_ptr<float>(), f2w.data_ptr<float>(), f2b.data_ptr<float>(), y4.data_ptr<float>(), B, 120, 84, 1);
    linear_tiled_kernel<<<dim3((C+TN-1)/TN, (B+TM-1)/TM), block, 0, stream>>>(y4.data_ptr<float>(), f3w.data_ptr<float>(), f3b.data_ptr<float>(), out.data_ptr<float>(), B, 84, C, 0);
    return out;
}
"""

_ext = load_inline(name="lenet5_fused_ext_final", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["lenet_forward"], extra_cuda_cflags=["-O3"], extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, num_classes):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        self.fc1 = nn.Linear(in_features=16*5*5, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=num_classes)
    def forward(self, x):
        return _ext.lenet_forward(x, self.conv1.weight, self.conv1.bias, self.conv2.weight, self.conv2.bias, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.fc3.weight, self.fc3.bias)