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_source=r'''
#include <torch/extension.h>
void fire_forward_cuda(torch::Tensor x, torch::Tensor sq_w, torch::Tensor sq_b, torch::Tensor e1_w, torch::Tensor e1_b, torch::Tensor e3_w, torch::Tensor e3_b, torch::Tensor tmp, torch::Tensor out);
torch::Tensor fire_forward(torch::Tensor x, torch::Tensor sq_w, torch::Tensor sq_b, torch::Tensor e1_w, torch::Tensor e1_b, torch::Tensor e3_w, torch::Tensor e3_b){auto N=x.size(0),H=x.size(2),W=x.size(3);auto tmp=torch::empty({N,sq_w.size(0),H,W},x.options());auto out=torch::empty({N,e1_w.size(0)+e3_w.size(0),H,W},x.options());fire_forward_cuda(x,sq_w,sq_b,e1_w,e1_b,e3_w,e3_b,tmp,out);return out;}
'''
_cuda_source=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define BLOCK 256
#define S 6
#define C 3
#define E 64
#define G 8
__constant__ float c_sq_w[S*C]; __constant__ float c_sq_b[S]; __constant__ float c_e1_w[E*S]; __constant__ float c_e1_b[E]; __constant__ float c_e3_w[E*S*9]; __constant__ float c_e3_b[E];
__global__ void squeeze_kernel_c(const float* __restrict__ x,float* __restrict__ tmp,int N,int H,int W){
 long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; long long total=(long long)N*H*W; if(idx>=total)return; int hw=H*W; int pos=idx%hw; int n=idx/hw; const float* xp=x+n*C*hw+pos; float x0=xp[0],x1=xp[hw],x2=xp[2*hw]; float* tp=tmp+n*S*hw+pos;
 #pragma unroll
 for(int s=0;s<S;++s){float a=c_sq_b[s]; a=fmaf(x0,c_sq_w[s*C+0],a); a=fmaf(x1,c_sq_w[s*C+1],a); a=fmaf(x2,c_sq_w[s*C+2],a); tp[s*hw]=a>0.f?a:0.f;}
}
__global__ void expand1_g8(const float* __restrict__ tmp,float* __restrict__ out,int N,int H,int W){
 int hw=H*W; long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; long long total=(long long)N*(E/G)*hw; if(idx>=total)return; int pos=idx%hw; long long t=idx/hw; int g=t%(E/G); int n=t/(E/G); int oc=g*G; const float* tp=tmp+n*S*hw+pos;
 float a0=c_e1_b[oc+0],a1=c_e1_b[oc+1],a2=c_e1_b[oc+2],a3=c_e1_b[oc+3],a4=c_e1_b[oc+4],a5=c_e1_b[oc+5],a6=c_e1_b[oc+6],a7=c_e1_b[oc+7];
 #pragma unroll
 for(int s=0;s<S;++s){float v=tp[s*hw]; a0=fmaf(v,c_e1_w[(oc+0)*S+s],a0); a1=fmaf(v,c_e1_w[(oc+1)*S+s],a1); a2=fmaf(v,c_e1_w[(oc+2)*S+s],a2); a3=fmaf(v,c_e1_w[(oc+3)*S+s],a3); a4=fmaf(v,c_e1_w[(oc+4)*S+s],a4); a5=fmaf(v,c_e1_w[(oc+5)*S+s],a5); a6=fmaf(v,c_e1_w[(oc+6)*S+s],a6); a7=fmaf(v,c_e1_w[(oc+7)*S+s],a7);}
 float* op=out+n*(2*E)*hw+oc*hw+pos; op[0*hw]=a0>0.f?a0:0.f; op[1*hw]=a1>0.f?a1:0.f; op[2*hw]=a2>0.f?a2:0.f; op[3*hw]=a3>0.f?a3:0.f; op[4*hw]=a4>0.f?a4:0.f; op[5*hw]=a5>0.f?a5:0.f; op[6*hw]=a6>0.f?a6:0.f; op[7*hw]=a7>0.f?a7:0.f;
}
__global__ void expand3_g8(const float* __restrict__ tmp,float* __restrict__ out,int N,int H,int Wd){
 int hw=H*Wd; long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; long long total=(long long)N*(E/G)*hw; if(idx>=total)return; int pos=idx%hw; int y=pos/Wd; int x=pos-y*Wd; long long t=idx/hw; int g=t%(E/G); int n=t/(E/G); int oc=g*G;
 float a0=c_e3_b[oc+0],a1=c_e3_b[oc+1],a2=c_e3_b[oc+2],a3=c_e3_b[oc+3],a4=c_e3_b[oc+4],a5=c_e3_b[oc+5],a6=c_e3_b[oc+6],a7=c_e3_b[oc+7]; const float* base=tmp+n*S*hw;
 #pragma unroll
 for(int s=0;s<S;++s){const float* sp=base+s*hw;
  #pragma unroll
  for(int ky=0;ky<3;++ky){int yy=y+ky-1; if((unsigned)yy<(unsigned)H){int row=yy*Wd;
   #pragma unroll
   for(int kx=0;kx<3;++kx){int xx=x+kx-1; if((unsigned)xx<(unsigned)Wd){float v=sp[row+xx]; int kk=s*9+ky*3+kx; a0=fmaf(v,c_e3_w[(oc+0)*S*9+kk],a0); a1=fmaf(v,c_e3_w[(oc+1)*S*9+kk],a1); a2=fmaf(v,c_e3_w[(oc+2)*S*9+kk],a2); a3=fmaf(v,c_e3_w[(oc+3)*S*9+kk],a3); a4=fmaf(v,c_e3_w[(oc+4)*S*9+kk],a4); a5=fmaf(v,c_e3_w[(oc+5)*S*9+kk],a5); a6=fmaf(v,c_e3_w[(oc+6)*S*9+kk],a6); a7=fmaf(v,c_e3_w[(oc+7)*S*9+kk],a7);}}
  }}
 }
 float* op=out+n*(2*E)*hw+(E+oc)*hw+pos; op[0*hw]=a0>0.f?a0:0.f; op[1*hw]=a1>0.f?a1:0.f; op[2*hw]=a2>0.f?a2:0.f; op[3*hw]=a3>0.f?a3:0.f; op[4*hw]=a4>0.f?a4:0.f; op[5*hw]=a5>0.f?a5:0.f; op[6*hw]=a6>0.f?a6:0.f; op[7*hw]=a7>0.f?a7:0.f;
}
void fire_forward_cuda(torch::Tensor x,torch::Tensor sq_w,torch::Tensor sq_b,torch::Tensor e1_w,torch::Tensor e1_b,torch::Tensor e3_w,torch::Tensor e3_b,torch::Tensor tmp,torch::Tensor out){int N=x.size(0),H=x.size(2),W=x.size(3); cudaStream_t stream=at::cuda::getCurrentCUDAStream(); cudaMemcpyToSymbolAsync(c_sq_w,sq_w.data_ptr<float>(),S*C*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); cudaMemcpyToSymbolAsync(c_sq_b,sq_b.data_ptr<float>(),S*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); cudaMemcpyToSymbolAsync(c_e1_w,e1_w.data_ptr<float>(),E*S*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); cudaMemcpyToSymbolAsync(c_e1_b,e1_b.data_ptr<float>(),E*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); cudaMemcpyToSymbolAsync(c_e3_w,e3_w.data_ptr<float>(),E*S*9*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); cudaMemcpyToSymbolAsync(c_e3_b,e3_b.data_ptr<float>(),E*sizeof(float),0,cudaMemcpyDeviceToDevice,stream); long long nhw=(long long)N*H*W; squeeze_kernel_c<<<(nhw+BLOCK-1)/BLOCK,BLOCK,0,stream>>>(x.data_ptr<float>(),tmp.data_ptr<float>(),N,H,W); long long tot=(long long)N*(E/G)*H*W; expand1_g8<<<(tot+BLOCK-1)/BLOCK,BLOCK,0,stream>>>(tmp.data_ptr<float>(),out.data_ptr<float>(),N,H,W); expand3_g8<<<(tot+BLOCK-1)/BLOCK,BLOCK,0,stream>>>(tmp.data_ptr<float>(),out.data_ptr<float>(),N,H,W);}
'''
_ext=load_inline(name='fire_ext_final_g8',cpp_sources=[_cpp_source],cuda_sources=[_cuda_source],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 _ext.fire_forward(x,self.squeeze.weight,self.squeeze.bias,self.expand1x1.weight,self.expand1x1.bias,self.expand3x3.weight,self.expand3x3.bias)