import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, applies HardSwish, and then ReLU.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height, width).
        """
        x = self.conv(x)
        x = torch.nn.functional.hardswish(x)
        x = torch.relu(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
kernel_size = 3

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]

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

cpp_src=r"""
#include <torch/extension.h>
torch::Tensor conv_hswish_relu_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b);
"""

cuda_src=r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float actv(float a){
    return (a<=0.0f)?0.0f:((a>=3.0f)?a:a*(a+3.0f)*(1.0f/6.0f));
}

__global__ __launch_bounds__(256,1) void conv3_oc8_x2_kernel(
    const float* __restrict__ x, const float* __restrict__ w,
    const float* __restrict__ b, float* __restrict__ y, int N, int H, int W){
    int Ho=H-2, Wo=W-2, Wop=(Wo+1)>>1, spatial=Ho*Wo;
    int total=N*8*Ho*Wop;
    int idx=blockIdx.x*blockDim.x+threadIdx.x;
    if(idx>=total) return;
    int p=idx%(Ho*Wop);
    int gt=idx/(Ho*Wop);
    int g=gt&7;
    int n=gt>>3;
    int oh=p/Wop;
    int pw=p-oh*Wop;
    int ow=pw<<1;
    int ocb=g<<3;

    float a0=b[ocb],a1=b[ocb+1],a2=b[ocb+2],a3=b[ocb+3],a4=b[ocb+4],a5=b[ocb+5],a6=b[ocb+6],a7=b[ocb+7];
    float c0=a0,c1=a1,c2=a2,c3=a3,c4=a4,c5=a5,c6=a6,c7=a7;
    const float* xp0=x+((n*8*H+oh)*W+ow);
    const float* wp=w+ocb*72;

    #pragma unroll
    for(int ic=0; ic<8; ++ic){
        const float* xi=xp0+ic*H*W;
        int kk=ic*9;
        float xa, xb;
#define DO(off,xoff) xa=xi[xoff]; xb=xi[(xoff)+1]; \
        a0=fmaf(xa,wp[0*72+kk+off],a0); c0=fmaf(xb,wp[0*72+kk+off],c0); \
        a1=fmaf(xa,wp[1*72+kk+off],a1); c1=fmaf(xb,wp[1*72+kk+off],c1); \
        a2=fmaf(xa,wp[2*72+kk+off],a2); c2=fmaf(xb,wp[2*72+kk+off],c2); \
        a3=fmaf(xa,wp[3*72+kk+off],a3); c3=fmaf(xb,wp[3*72+kk+off],c3); \
        a4=fmaf(xa,wp[4*72+kk+off],a4); c4=fmaf(xb,wp[4*72+kk+off],c4); \
        a5=fmaf(xa,wp[5*72+kk+off],a5); c5=fmaf(xb,wp[5*72+kk+off],c5); \
        a6=fmaf(xa,wp[6*72+kk+off],a6); c6=fmaf(xb,wp[6*72+kk+off],c6); \
        a7=fmaf(xa,wp[7*72+kk+off],a7); c7=fmaf(xb,wp[7*72+kk+off],c7);
        DO(0,0); DO(1,1); DO(2,2);
        DO(3,W); DO(4,W+1); DO(5,W+2);
        DO(6,2*W); DO(7,2*W+1); DO(8,2*W+2);
#undef DO
    }
    int s=oh*Wo+ow;
    int ybase=(n*64+ocb)*spatial+s;
    y[ybase]=actv(a0); y[ybase+spatial]=actv(a1); y[ybase+2*spatial]=actv(a2); y[ybase+3*spatial]=actv(a3);
    y[ybase+4*spatial]=actv(a4); y[ybase+5*spatial]=actv(a5); y[ybase+6*spatial]=actv(a6); y[ybase+7*spatial]=actv(a7);
    if(ow+1<Wo){
        y[ybase+1]=actv(c0); y[ybase+spatial+1]=actv(c1); y[ybase+2*spatial+1]=actv(c2); y[ybase+3*spatial+1]=actv(c3);
        y[ybase+4*spatial+1]=actv(c4); y[ybase+5*spatial+1]=actv(c5); y[ybase+6*spatial+1]=actv(c6); y[ybase+7*spatial+1]=actv(c7);
    }
}

__global__ void conv_generic(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,int N,int C,int OC,int H,int W,int K){
    int Ho=H-K+1,Wo=W-K+1;
    long total=(long)N*OC*Ho*Wo;
    long idx=(long)blockIdx.x*blockDim.x+threadIdx.x;
    if(idx>=total) return;
    int ow=idx%Wo;
    long t=idx/Wo;
    int oh=t%Ho;
    t/=Ho;
    int oc=t%OC;
    int n=t/OC;
    float acc=b[oc];
    for(int ic=0; ic<C; ++ic)
        for(int kh=0; kh<K; ++kh)
            for(int kw=0; kw<K; ++kw)
                acc=fmaf(x[((long)n*C+ic)*H*W+(oh+kh)*W+(ow+kw)], w[((long)oc*C+ic)*K*K+kh*K+kw], acc);
    y[idx]=actv(acc);
}

torch::Tensor conv_hswish_relu_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b){
    int N=x.size(0), C=x.size(1), H=x.size(2), W=x.size(3);
    int OC=w.size(0), K=w.size(2);
    int Ho=H-K+1, Wo=W-K+1;
    auto y=torch::empty({N,OC,Ho,Wo}, x.options());
    int threads=256;
    if(C==8 && OC==64 && K==3){
        int total=N*8*Ho*((Wo+1)>>1);
        int blocks=(total+threads-1)/threads;
        conv3_oc8_x2_kernel<<<blocks,threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), y.data_ptr<float>(), N, H, W);
    }else{
        long total=(long)N*OC*Ho*Wo;
        int blocks=(total+threads-1)/threads;
        conv_generic<<<blocks,threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), y.data_ptr<float>(), N, C, OC, H, W, K);
    }
    return y;
}
"""

_ext=load_inline(name="conv_hswish_relu_ext_final", cpp_sources=cpp_src, cuda_sources=cuda_src,
                 functions=["conv_hswish_relu_forward"], verbose=False,
                 extra_cuda_cflags=["-O3", "--use_fast_math"])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.conv=nn.Conv2d(in_channels, out_channels, kernel_size)
    def forward(self, x):
        return _ext.conv_hswish_relu_forward(x, self.conv.weight, self.conv.bias)