import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, pool_stride, pool_padding):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.max_pool = nn.MaxPool3d(kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding)
        self.subtract = nn.Parameter(torch.randn(out_channels)) # Assuming subtraction is element-wise across channels

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.max_pool(x)
        x = torch.softmax(x, dim=1) # Apply softmax across channels (dim=1)
        x = x - self.subtract.view(1, -1, 1, 1, 1) # Subtract across channels
        x = torch.sigmoid(x) * x # Swish activation
        x = torch.max(x, dim=1)[0] # Max pooling across channels
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
pool_kernel_size = 2
pool_stride = 2
pool_padding = 0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, pool_stride, pool_padding]

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 weight, torch::Tensor bias, torch::Tensor sub);
"""
cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
#define IC 3
#define OC 16
#define K 3
__device__ __forceinline__ float pooled_oc_direct(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,int n,int oc,int pd,int ph,int pw,int D,int H,int W){
    float v0=b[oc], v1=v0, v2=v0, v3=v0, v4=v0, v5=v0, v6=v0, v7=v0;
    #pragma unroll
    for(int ic=0; ic<IC; ++ic){
        #pragma unroll
        for(int rd=0; rd<2; ++rd){
            int id = pd + rd;
            if(id < D){
                #pragma unroll
                for(int rh=0; rh<2; ++rh){
                    int ih = ph + rh;
                    if(ih < H){
                        #pragma unroll
                        for(int rw=0; rw<2; ++rw){
                            int iw = pw + rw;
                            if(iw < W){
                                float xv = x[(((long)n*IC + ic)*D + id)*H*W + ih*W + iw];
                                int dd0 = rd; int dh0 = rh; int dw0 = rw;
                                #pragma unroll
                                for(int dd=0; dd<2; ++dd){
                                    if(dd >= dd0){
                                        int kd = (dd==0 ? 1 : (rd==0 ? 2 : 0));
                                        #pragma unroll
                                        for(int dh=0; dh<2; ++dh){
                                            if(dh >= dh0){
                                                int kh = (dh==0 ? 1 : (rh==0 ? 2 : 0));
                                                #pragma unroll
                                                for(int dw=0; dw<2; ++dw){
                                                    if(dw >= dw0){
                                                        int kw = (dw==0 ? 1 : (rw==0 ? 2 : 0));
                                                        float ww = w[((((long)ic*OC + oc)*K + kd)*K + kh)*K + kw];
                                                        int p = (dd<<2) | (dh<<1) | dw;
                                                        if(p==0) v0 = fmaf(xv, ww, v0);
                                                        else if(p==1) v1 = fmaf(xv, ww, v1);
                                                        else if(p==2) v2 = fmaf(xv, ww, v2);
                                                        else if(p==3) v3 = fmaf(xv, ww, v3);
                                                        else if(p==4) v4 = fmaf(xv, ww, v4);
                                                        else if(p==5) v5 = fmaf(xv, ww, v5);
                                                        else if(p==6) v6 = fmaf(xv, ww, v6);
                                                        else v7 = fmaf(xv, ww, v7);
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return fmaxf(fmaxf(fmaxf(v0,v1),fmaxf(v2,v3)), fmaxf(fmaxf(v4,v5),fmaxf(v6,v7)));
}
__global__ void fused_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ sub,float* __restrict__ out,int N,int D,int H,int W){
    long idx=(long)blockIdx.x*blockDim.x+threadIdx.x; long total=(long)N*D*H*W; if(idx>=total)return;
    int pw=idx%W; long t=idx/W; int ph=t%H; t/=H; int pd=t%D; int n=t/D;
    float m[OC]; float maxm=-FLT_MAX;
    #pragma unroll
    for(int oc=0; oc<OC; ++oc){ float mv=pooled_oc_direct(x,w,b,n,oc,pd,ph,pw,D,H,W); m[oc]=mv; maxm=fmaxf(maxm,mv); }
    float denom=0.0f;
    #pragma unroll
    for(int oc=0; oc<OC; ++oc){ float e=expf(m[oc]-maxm); m[oc]=e; denom += e; }
    float inv=1.0f/denom; float best=-FLT_MAX;
    #pragma unroll
    for(int oc=0; oc<OC; ++oc){ float z=m[oc]*inv-sub[oc]; float y=z/(1.0f+expf(-z)); best=fmaxf(best,y); }
    out[idx]=best;
}
torch::Tensor fused_forward(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor sub){int N=x.size(0),D=x.size(2),H=x.size(3),W=x.size(4);auto out=torch::empty({N,D,H,W},x.options());long total=(long)N*D*H*W;int threads=256;int blocks=(total+threads-1)/threads;fused_kernel<<<blocks,threads>>>(x.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),sub.data_ptr<float>(),out.data_ptr<float>(),N,D,H,W);return out;}
"""
_ext=load_inline(name="fused_deconv_pool_direct_t256_final",cpp_sources=cpp_src,cuda_sources=cuda_src,functions=["fused_forward"],extra_cuda_cflags=["-O3"],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride,padding,output_padding,pool_kernel_size,pool_stride,pool_padding):
        super().__init__(); self.conv_transpose=nn.ConvTranspose3d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,output_padding=output_padding); self.max_pool=nn.MaxPool3d(kernel_size=pool_kernel_size,stride=pool_stride,padding=pool_padding); self.subtract=nn.Parameter(torch.randn(out_channels))
    def forward(self,x):
        return _ext.fused_forward(x.contiguous(),self.conv_transpose.weight.contiguous(),self.conv_transpose.bias.contiguous(),self.subtract.contiguous())