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
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAEvent.h>
#include <optional>

using torch::Tensor;

static inline Tensor conv2d_wrap(const Tensor& input, const Tensor& weight,
                                 const Tensor& bias, int64_t pad) {
    std::optional<Tensor> obias(bias);
    return at::conv2d(input, weight, obias,
                      c10::IntArrayRef({1,1}),
                      c10::IntArrayRef({pad,pad}),
                      c10::IntArrayRef({1,1}), 1);
}

Tensor inception_forward(const Tensor& x,
                         const Tensor& w1,  const Tensor& b1,
                         const Tensor& wr3, const Tensor& br3,
                         const Tensor& w3,  const Tensor& b3,
                         const Tensor& wr5, const Tensor& br5,
                         const Tensor& w5,  const Tensor& b5,
                         const Tensor& wp,  const Tensor& bp) {
    Tensor y1, red, r3, r5, y3, y5, pooled, yp;
    const int64_t c = x.size(1);

    // The two reduction 1x1 convolutions are independent and have identical
    // inputs.  Execute them as one wider 1x1 convolution; each output channel
    // is computed identically to the separate convolutions, but it removes one
    // large input reread and one cuDNN launch.
    Tensor wcomb = at::cat({wr3.view({wr3.size(0), c}),
                            wr5.view({wr5.size(0), c})}, 0)
                       .view({wr3.size(0) + wr5.size(0), c, 1, 1});
    Tensor bcomb = at::cat({br3, br5}, 0);

    const int dev = x.device().index();
    auto s1 = at::cuda::getStreamFromPool(false, dev);
    auto s2 = at::cuda::getStreamFromPool(false, dev);
    auto s3 = at::cuda::getStreamFromPool(false, dev);
    auto s4 = at::cuda::getStreamFromPool(false, dev);
    auto cur = at::cuda::getCurrentCUDAStream(dev);

    c10::cuda::CUDAEvent e1, ered, e2, e3, e4;

    // Branches are independent, so schedule them on separate streams.  The
    // 3x3/5x5 branches share the combined reduction; once that finishes their
    // second-stage convolutions can overlap.
    {
        c10::cuda::CUDAStreamGuard guard(s1);
        y1 = conv2d_wrap(x, w1, b1, 0);
        e1.record(s1);
    }
    {
        c10::cuda::CUDAStreamGuard guard(s2);
        red = conv2d_wrap(x, wcomb, bcomb, 0);
        r3 = red.narrow(1, 0, wr3.size(0));
        r5 = red.narrow(1, wr3.size(0), wr5.size(0));
        ered.record(s2);
        y3 = conv2d_wrap(r3, w3, b3, 1);
        e2.record(s2);
    }
    {
        c10::cuda::CUDAStreamGuard guard(s3);
        ered.block(s3);
        y5 = conv2d_wrap(r5, w5, b5, 2);
        e3.record(s3);
    }
    {
        c10::cuda::CUDAStreamGuard guard(s4);
        pooled = at::max_pool2d(x,
                                c10::IntArrayRef({3,3}),
                                c10::IntArrayRef({1,1}),
                                c10::IntArrayRef({1,1}),
                                c10::IntArrayRef({1,1}), false);
        yp = conv2d_wrap(pooled, wp, bp, 0);
        e4.record(s4);
    }

    e1.block(cur);
    e2.block(cur);
    e3.block(cur);
    e4.block(cur);
    return at::cat({y1, y3, y5, yp}, 1);
}
'''

# Supplying a CUDA translation unit makes the CUDA headers/libraries available
# to the C++ TU that manages streams/events.
_cuda_src = r'''
__global__ void _inception_stream_dummy_kernel() {}
'''

_ext = load_inline(
    name='inception_streamed_combined_ext_submit',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['inception_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    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)
        )

    def forward(self, x):
        return _ext.inception_forward(
            x,
            self.branch1x1.weight, self.branch1x1.bias,
            self.branch3x3[0].weight, self.branch3x3[0].bias,
            self.branch3x3[1].weight, self.branch3x3[1].bias,
            self.branch5x5[0].weight, self.branch5x5[0].bias,
            self.branch5x5[1].weight, self.branch5x5[1].bias,
            self.branch_pool[1].weight, self.branch_pool[1].bias,
        )