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

class Model(nn.Module):
    def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
        """
        :param in_channels: Number of input channels
        :param out_1x1: Number of output channels for the 1x1 convolution
        :param reduce_3x3: Number of output channels for the 1x1 reduction before 3x3 convolution
        :param out_3x3: Number of output channels for the 3x3 convolution
        :param reduce_5x5: Number of output channels for the 1x1 reduction before 5x5 convolution
        :param out_5x5: Number of output channels for the 5x5 convolution
        :param pool_proj: Number of output channels for the pooling projection
        """
        super(Model, self).__init__()
        
        # 1x1 convolution branch
        self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
        
        # 3x3 convolution branch
        self.branch3x3 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
            nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1)
        )
        
        # 5x5 convolution branch
        self.branch5x5 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
            nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2)
        )
        
        # Max pooling branch
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1)
        )
    
    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels, height, width)
        """
        branch1x1 = self.branch1x1(x)
        branch3x3 = self.branch3x3(x)
        branch5x5 = self.branch5x5(x)
        branch_pool = self.branch_pool(x)
        
        outputs = [branch1x1, branch3x3, branch5x5, branch_pool]
        return torch.cat(outputs, 1)

# Test code
in_channels = 480
out_1x1 = 192
reduce_3x3 = 96
out_3x3 = 208
reduce_5x5 = 16
out_5x5 = 48
pool_proj = 64
batch_size = 10
height = 224
width = 224

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj]

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 maxpool3x3_cuda(torch::Tensor x);
"""

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

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

    int w = (int)(idx % W);
    long t = idx / W;
    int h = (int)(t % H);
    long base = (t / H) * (long)H * (long)W;

    const float* p = x + base;
    float m;

    if (h > 0 && h < H - 1 && w > 0 && w < W - 1) {
        long c = (long)h * W + w;
        m = p[c - W - 1];
        float v;
        v = p[c - W    ]; m = fmaxf(m, v);
        v = p[c - W + 1]; m = fmaxf(m, v);
        v = p[c     - 1]; m = fmaxf(m, v);
        v = p[c        ]; m = fmaxf(m, v);
        v = p[c     + 1]; m = fmaxf(m, v);
        v = p[c + W - 1]; m = fmaxf(m, v);
        v = p[c + W    ]; m = fmaxf(m, v);
        v = p[c + W + 1]; m = fmaxf(m, v);
    } else {
        m = -FLT_MAX;
        for (int dh = -1; dh <= 1; ++dh) {
            int hh = h + dh;
            if ((unsigned)hh >= (unsigned)H) continue;
            for (int dw = -1; dw <= 1; ++dw) {
                int ww = w + dw;
                if ((unsigned)ww >= (unsigned)W) continue;
                float v = p[(long)hh * W + ww];
                m = fmaxf(m, v);
            }
        }
    }
    y[idx] = m;
}

torch::Tensor maxpool3x3_cuda(torch::Tensor x) {
    auto y = torch::empty_like(x);
    const long total = x.numel();
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    maxpool3x3_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), total, H, W);
    return y;
}
"""

_ext = load_inline(
    name="inception_pool3x3_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["maxpool3x3_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
        super(ModelNew, self).__init__()
        self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
        self.branch3x3 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
            nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1)
        )
        self.branch5x5 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
            nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2)
        )
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1)
        )
        self._out_1x1 = out_1x1
        self._reduce_3x3 = reduce_3x3
        self._reduce_5x5 = reduce_5x5

    def forward(self, x):
        # Fuse the three independent 1x1 convolutions into a single cuDNN call.
        w0 = torch.cat((self.branch1x1.weight,
                        self.branch3x3[0].weight,
                        self.branch5x5[0].weight), dim=0)
        b0 = torch.cat((self.branch1x1.bias,
                        self.branch3x3[0].bias,
                        self.branch5x5[0].bias), dim=0)
        y0 = F.conv2d(x, w0, b0)

        c1 = self._out_1x1
        c3 = self._reduce_3x3
        c5 = self._reduce_5x5
        branch1x1 = y0[:, :c1, :, :]
        red3 = y0[:, c1:c1 + c3, :, :]
        red5 = y0[:, c1 + c3:c1 + c3 + c5, :, :]

        branch3x3 = F.conv2d(red3, self.branch3x3[1].weight, self.branch3x3[1].bias, padding=1)
        branch5x5 = F.conv2d(red5, self.branch5x5[1].weight, self.branch5x5[1].bias, padding=2)

        pooled = _ext.maxpool3x3_cuda(x)
        branch_pool = F.conv2d(pooled, self.branch_pool[1].weight, self.branch_pool[1].bias)

        return torch.cat((branch1x1, branch3x3, branch5x5, branch_pool), dim=1)