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 w1, torch::Tensor b1,
                           torch::Tensor w3, torch::Tensor b3);
"""

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

#define C_IN 3
#define S_CH 6
#define E1_CH 64
#define E3_CH 64
#define E_CH 128
#define TW 16
#define TH 8
#define OG 16
#define PIX (TW*TH)
#define SHW (TW+2)
#define SHH (TH+2)

__global__ void squeeze_kernel(const float* __restrict__ x,
                               const float* __restrict__ w,
                               const float* __restrict__ b,
                               float* __restrict__ s,
                               long long total_pix, int H, int W) {
    long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= total_pix) return;
    int HW = H * W;
    int p = (int)(t % HW);
    long long n = t / HW;
    const float* xp = x + n * (long long)C_IN * HW + p;
    float x0 = xp[0LL * HW];
    float x1 = xp[1LL * HW];
    float x2 = xp[2LL * HW];
    float* sp = s + n * (long long)S_CH * HW + p;
#pragma unroll
    for (int sc = 0; sc < S_CH; ++sc) {
        float v = b[sc];
        v = fmaf(x0, w[sc*C_IN + 0], v);
        v = fmaf(x1, w[sc*C_IN + 1], v);
        v = fmaf(x2, w[sc*C_IN + 2], v);
        sp[(long long)sc * HW] = v > 0.0f ? v : 0.0f;
    }
}

__global__ void expand1_tiled_kernel(const float* __restrict__ s,
                                     const float* __restrict__ w1,
                                     const float* __restrict__ b1,
                                     float* __restrict__ out,
                                     int N, int H, int W) {
    __shared__ float sh[S_CH * PIX];
    int HW = H * W;
    int tile_x = blockIdx.x * TW;
    int tile_y = blockIdx.y * TH;
    int groups_per_n = E1_CH / OG;
    int n = blockIdx.z / groups_per_n;
    int g = blockIdx.z - n * groups_per_n;

    for (int t = threadIdx.x; t < S_CH * PIX; t += blockDim.x) {
        int pix = t % PIX;
        int ic = t / PIX;
        int py = pix / TW;
        int px = pix - py * TW;
        int gy = tile_y + py;
        int gx = tile_x + px;
        float v = 0.0f;
        if (gy < H && gx < W) {
            v = s[((long long)n * S_CH + ic) * HW + gy * W + gx];
        }
        sh[t] = v;
    }
    __syncthreads();

    for (int idx = threadIdx.x; idx < OG * PIX; idx += blockDim.x) {
        int pix = idx % PIX;
        int ocl = idx / PIX;
        int py = pix / TW;
        int px = pix - py * TW;
        int gy = tile_y + py;
        int gx = tile_x + px;
        int oc = g * OG + ocl;
        if (gy < H && gx < W) {
            float sum = b1[oc];
#pragma unroll
            for (int ic = 0; ic < S_CH; ++ic) {
                sum = fmaf(sh[ic * PIX + pix], w1[oc * S_CH + ic], sum);
            }
            out[((long long)n * E_CH + oc) * HW + gy * W + gx] = sum > 0.0f ? sum : 0.0f;
        }
    }
}

__global__ void expand3_tiled_kernel(const float* __restrict__ s,
                                     const float* __restrict__ w3,
                                     const float* __restrict__ b3,
                                     float* __restrict__ out,
                                     int N, int H, int W) {
    __shared__ float sh[S_CH * SHH * SHW];
    int HW = H * W;
    int tile_x = blockIdx.x * TW;
    int tile_y = blockIdx.y * TH;
    int groups_per_n = E3_CH / OG;
    int n = blockIdx.z / groups_per_n;
    int g = blockIdx.z - n * groups_per_n;

    for (int t = threadIdx.x; t < S_CH * SHH * SHW; t += blockDim.x) {
        int rem = t;
        int sx = rem % SHW; rem /= SHW;
        int sy = rem % SHH;
        int ic = rem / SHH;
        int gy = tile_y + sy - 1;
        int gx = tile_x + sx - 1;
        float v = 0.0f;
        if ((unsigned)gy < (unsigned)H && (unsigned)gx < (unsigned)W) {
            v = s[((long long)n * S_CH + ic) * HW + gy * W + gx];
        }
        sh[t] = v;
    }
    __syncthreads();

    for (int idx = threadIdx.x; idx < OG * PIX; idx += blockDim.x) {
        int pix = idx % PIX;
        int ocl = idx / PIX;
        int py = pix / TW;
        int px = pix - py * TW;
        int gy = tile_y + py;
        int gx = tile_x + px;
        int oc = g * OG + ocl;
        if (gy < H && gx < W) {
            float sum = b3[oc];
#pragma unroll
            for (int ic = 0; ic < S_CH; ++ic) {
                const float* shp = sh + ic * SHH * SHW + py * SHW + px;
                const float* wp = w3 + (oc * S_CH + ic) * 9;
                sum = fmaf(shp[0 * SHW + 0], wp[0], sum);
                sum = fmaf(shp[0 * SHW + 1], wp[1], sum);
                sum = fmaf(shp[0 * SHW + 2], wp[2], sum);
                sum = fmaf(shp[1 * SHW + 0], wp[3], sum);
                sum = fmaf(shp[1 * SHW + 1], wp[4], sum);
                sum = fmaf(shp[1 * SHW + 2], wp[5], sum);
                sum = fmaf(shp[2 * SHW + 0], wp[6], sum);
                sum = fmaf(shp[2 * SHW + 1], wp[7], sum);
                sum = fmaf(shp[2 * SHW + 2], wp[8], sum);
            }
            out[((long long)n * E_CH + (E1_CH + oc)) * HW + gy * W + gx] = sum > 0.0f ? sum : 0.0f;
        }
    }
}

torch::Tensor fire_forward(torch::Tensor x,
                           torch::Tensor sw, torch::Tensor sb,
                           torch::Tensor w1, torch::Tensor b1,
                           torch::Tensor w3, torch::Tensor b3) {
    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    auto s = torch::empty({N, S_CH, H, W}, x.options());
    auto out = torch::empty({N, E_CH, H, W}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    long long total_pix = (long long)N * H * W;
    int threads = 256;
    int blocks = (int)((total_pix + threads - 1) / threads);
    squeeze_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), sw.data_ptr<float>(), sb.data_ptr<float>(),
                                                   s.data_ptr<float>(), total_pix, H, W);

    dim3 block2(256);
    dim3 grid((W + TW - 1) / TW, (H + TH - 1) / TH, N * (E1_CH / OG));
    expand1_tiled_kernel<<<grid, block2, 0, stream>>>(s.data_ptr<float>(), w1.data_ptr<float>(), b1.data_ptr<float>(),
                                                      out.data_ptr<float>(), N, H, W);
    expand3_tiled_kernel<<<grid, block2, 0, stream>>>(s.data_ptr<float>(), w3.data_ptr<float>(), b3.data_ptr<float>(),
                                                      out.data_ptr<float>(), N, H, W);
    return out;
}
"""

_fire_ext = load_inline(
    name="fire_module_fused_tiled_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["fire_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels):
        super(ModelNew, 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):
        return _fire_ext.fire_forward(x, self.squeeze.weight, self.squeeze.bias,
                                      self.expand1x1.weight, self.expand1x1.bias,
                                      self.expand3x3.weight, self.expand3x3.bias)