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

class Model(nn.Module):
    def __init__(self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels):
        """
        :param in_channels: Number of input channels
        :param squeeze_channels: Number of output channels for the squeeze layer
        :param expand1x1_channels: Number of output channels for the 1x1 expand layer
        :param expand3x3_channels: Number of output channels for the 3x3 expand layer
        """
        super(Model, self).__init__()
        
        self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
        self.squeeze_activation = nn.ReLU(inplace=True)
        
        self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
        self.expand1x1_activation = nn.ReLU(inplace=True)
        
        self.expand3x3 = nn.Conv2d(squeeze_channels, expand3x3_channels, kernel_size=3, padding=1)
        self.expand3x3_activation = nn.ReLU(inplace=True)
    
    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, expand1x1_channels + expand3x3_channels, height, width)
        """
        x = self.squeeze_activation(self.squeeze(x))
        return torch.cat([
            self.expand1x1_activation(self.expand1x1(x)),
            self.expand3x3_activation(self.expand3x3(x))
        ], 1)

# Test code
batch_size = 128
num_input_features = 3
num_output_features = 64
height, width = 256, 256
squeeze_channels = 6
expand1x1_channels = 64
expand3x3_channels = 64

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

def get_init_inputs():
    return [num_input_features, squeeze_channels, expand1x1_channels, expand3x3_channels]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fire_forward(torch::Tensor x, torch::Tensor sw, torch::Tensor sb, torch::Tensor ew1, torch::Tensor eb1, torch::Tensor ew3, torch::Tensor eb3);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define BLOCK_OC 8

__device__ __forceinline__ void compute_sq6_c3(const float* __restrict__ x,
    const float* __restrict__ sw, const float* __restrict__ sb,
    int n, int h, int w, int H, int W, float sq[6]) {
    int base = ((n * 3) * H + h) * W + w;
    float x0 = x[base];
    float x1 = x[base + H*W];
    float x2 = x[base + 2*H*W];
    #pragma unroll
    for (int s=0; s<6; ++s) {
        const float* ws = sw + s*3;
        float v = sb[s] + ws[0]*x0 + ws[1]*x1 + ws[2]*x2;
        sq[s] = v > 0.0f ? v : 0.0f;
    }
}

__global__ __launch_bounds__(256, 2) void fire_fused_kernel(
    const float* __restrict__ x, const float* __restrict__ sw, const float* __restrict__ sb,
    const float* __restrict__ ew1, const float* __restrict__ eb1,
    const float* __restrict__ ew3, const float* __restrict__ eb3,
    float* __restrict__ y, int N, int H, int W, int E1, int E3) {
    long pixels = (long)N * H * W;
    int maxE = E1 > E3 ? E1 : E3;
    int oc_blocks = (maxE + BLOCK_OC - 1) / BLOCK_OC;
    long total = pixels * oc_blocks;
    long tid = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    for (long idx = tid; idx < total; idx += stride) {
        long p = idx % pixels;
        int ob = (int)(idx / pixels);
        int w = (int)(p % W);
        int tmp = (int)(p / W);
        int h = tmp % H;
        int n = tmp / H;
        int oc0 = ob * BLOCK_OC;
        float sqc[6];
        compute_sq6_c3(x, sw, sb, n, h, w, H, W, sqc);
        float a1[BLOCK_OC];
        float a3[BLOCK_OC];
        #pragma unroll
        for (int i=0; i<BLOCK_OC; ++i) { int oc = oc0 + i; a1[i] = (oc < E1) ? eb1[oc] : 0.0f; a3[i] = (oc < E3) ? eb3[oc] : 0.0f; }
        #pragma unroll
        for (int s=0; s<6; ++s) { float sv = sqc[s];
            #pragma unroll
            for (int i=0; i<BLOCK_OC; ++i) { int oc = oc0 + i; if (oc < E1) a1[i] += ew1[oc*6 + s] * sv; }
        }
        #pragma unroll
        for (int ky=0; ky<3; ++ky) { int ih = h + ky - 1; if (ih < 0 || ih >= H) continue;
            #pragma unroll
            for (int kx=0; kx<3; ++kx) { int iw = w + kx - 1; if (iw < 0 || iw >= W) continue; float sq[6];
                if (ky == 1 && kx == 1) {
                    #pragma unroll
                    for (int s=0; s<6; ++s) sq[s] = sqc[s];
                } else { compute_sq6_c3(x, sw, sb, n, ih, iw, H, W, sq); }
                int kidx = ky*3 + kx;
                #pragma unroll
                for (int s=0; s<6; ++s) { float sv = sq[s];
                    #pragma unroll
                    for (int i=0; i<BLOCK_OC; ++i) { int oc = oc0 + i; if (oc < E3) a3[i] += ew3[(oc*6 + s)*9 + kidx] * sv; }
                }
            }
        }
        long spatial = (long)H * W;
        long out_base = ((long)n * (E1 + E3)) * spatial + (long)h * W + w;
        #pragma unroll
        for (int i=0; i<BLOCK_OC; ++i) { int oc = oc0 + i; if (oc < E1) { float v = a1[i]; y[out_base + (long)oc * spatial] = v > 0.0f ? v : 0.0f; } if (oc < E3) { float v = a3[i]; y[out_base + (long)(E1 + oc) * spatial] = v > 0.0f ? v : 0.0f; } }
    }
}

torch::Tensor fire_forward(torch::Tensor x, torch::Tensor sw, torch::Tensor sb, torch::Tensor ew1, torch::Tensor eb1, torch::Tensor ew3, torch::Tensor eb3) {
    int N = x.size(0), H = x.size(2), W = x.size(3), E1 = ew1.size(0), E3 = ew3.size(0);
    auto y = torch::empty({N, E1 + E3, H, W}, x.options());
    long pixels = (long)N * H * W; int maxE = E1 > E3 ? E1 : E3; int oc_blocks = (maxE + BLOCK_OC - 1) / BLOCK_OC; long total = pixels * oc_blocks;
    int threads = 256; int blocks = (int)((total + threads - 1) / threads); if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fire_fused_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), sw.data_ptr<float>(), sb.data_ptr<float>(), ew1.data_ptr<float>(), eb1.data_ptr<float>(), ew3.data_ptr<float>(), eb3.data_ptr<float>(), y.data_ptr<float>(), N, H, W, E1, E3);
    return y;
}
'''
_ext = load_inline(name='fire_fused_ext_submit', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['fire_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels):
        super().__init__()
        self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
        self.squeeze_activation = nn.ReLU(inplace=True)
        self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
        self.expand1x1_activation = nn.ReLU(inplace=True)
        self.expand3x3 = nn.Conv2d(squeeze_channels, expand3x3_channels, kernel_size=3, padding=1)
        self.expand3x3_activation = nn.ReLU(inplace=True)
    def forward(self, x):
        return _ext.fire_forward(x.contiguous(), self.squeeze.weight.contiguous(), self.squeeze.bias.contiguous(), self.expand1x1.weight.contiguous(), self.expand1x1.bias.contiguous(), self.expand3x3.weight.contiguous(), self.expand3x3.bias.contiguous())