import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a convolution, applies tanh, scaling, adds a bias term, and then max-pools.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape, pool_kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scaling_factor = scaling_factor
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.max_pool = nn.MaxPool2d(pool_kernel_size)

    def forward(self, x):
        # Convolution
        x = self.conv(x)
        # Tanh activation
        x = torch.tanh(x)
        # Scaling
        x = x * self.scaling_factor
        # Bias addition
        x = x + self.bias
        # Max-pooling
        x = self.max_pool(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 256, 256
kernel_size = 3
scaling_factor = 2.0
bias_shape = (out_channels, 1, 1)
pool_kernel_size = 4

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scaling_factor, bias_shape, pool_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 fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor cb, torch::Tensor addb, double scaling);
'''

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

__global__ void fused_conv_tanh_pool4_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ cb,const float* __restrict__ addb,float* __restrict__ y,int N,int H,int W,int OH,int OW,float scaling){
    long long idx=(long long)blockIdx.x*blockDim.x+threadIdx.x; long long total=(long long)N*64LL*OH*OW; if(idx>=total)return;
    int ow=idx%OW; int tmp=idx/OW; int oh=tmp%OH; tmp/=OH; int oc=tmp&63; int n=tmp>>6;
    int ih0=oh*4, iw0=ow*4; const float* xbase=x+((long long)n*8*H*W); const float* wbase=w+(long long)oc*72;
    float b0=cb[oc]; float acc[16];
    #pragma unroll
    for(int i=0;i<16;i++) acc[i]=b0;
    #pragma unroll
    for(int ic=0;ic<8;++ic){ const float* xp=xbase+(long long)ic*H*W+(long long)ih0*W+iw0; float p[36];
        #pragma unroll
        for(int rr=0;rr<6;++rr){
            #pragma unroll
            for(int cc=0;cc<6;++cc) p[rr*6+cc]=__ldg(xp+(long long)rr*W+cc);
        }
        const float* wp=wbase+ic*9;
        #pragma unroll
        for(int ky=0;ky<3;++ky){
            #pragma unroll
            for(int kx=0;kx<3;++kx){ float ww=__ldg(wp+ky*3+kx);
                acc[0]=fmaf(p[(ky+0)*6+(kx+0)],ww,acc[0]); acc[1]=fmaf(p[(ky+0)*6+(kx+1)],ww,acc[1]); acc[2]=fmaf(p[(ky+0)*6+(kx+2)],ww,acc[2]); acc[3]=fmaf(p[(ky+0)*6+(kx+3)],ww,acc[3]);
                acc[4]=fmaf(p[(ky+1)*6+(kx+0)],ww,acc[4]); acc[5]=fmaf(p[(ky+1)*6+(kx+1)],ww,acc[5]); acc[6]=fmaf(p[(ky+1)*6+(kx+2)],ww,acc[6]); acc[7]=fmaf(p[(ky+1)*6+(kx+3)],ww,acc[7]);
                acc[8]=fmaf(p[(ky+2)*6+(kx+0)],ww,acc[8]); acc[9]=fmaf(p[(ky+2)*6+(kx+1)],ww,acc[9]); acc[10]=fmaf(p[(ky+2)*6+(kx+2)],ww,acc[10]); acc[11]=fmaf(p[(ky+2)*6+(kx+3)],ww,acc[11]);
                acc[12]=fmaf(p[(ky+3)*6+(kx+0)],ww,acc[12]); acc[13]=fmaf(p[(ky+3)*6+(kx+1)],ww,acc[13]); acc[14]=fmaf(p[(ky+3)*6+(kx+2)],ww,acc[14]); acc[15]=fmaf(p[(ky+3)*6+(kx+3)],ww,acc[15]);
            }
        }
    }
    float m=acc[0]; if(scaling>=0.0f){
        #pragma unroll
        for(int i=1;i<16;i++) m=fmaxf(m,acc[i]);
    }else{
        #pragma unroll
        for(int i=1;i<16;i++) m=fminf(m,acc[i]);
    }
    y[idx]=tanhf(m)*scaling+__ldg(addb+oc);
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor cb, torch::Tensor addb, double scaling){
    int N=x.size(0), H=x.size(2), W=x.size(3); int OH=(H-6)/4+1, OW=(W-6)/4+1; auto y=torch::empty({N,64,OH,OW},x.options()); long long total=(long long)N*64LL*OH*OW; int threads=256; int blocks=(int)((total+threads-1)/threads);
    fused_conv_tanh_pool4_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),cb.data_ptr<float>(),addb.data_ptr<float>(),y.data_ptr<float>(),N,H,W,OH,OW,(float)scaling); return y;
}
'''

_ext = load_inline(name='fused_conv_tanh_pool4_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['fused_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape, pool_kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scaling_factor = scaling_factor
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.max_pool = nn.MaxPool2d(pool_kernel_size)
    def forward(self, x):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias, float(self.scaling_factor))