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 <cfloat>

__global__ void fused_deconv_pool_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    const float* __restrict__ sub,
    float* __restrict__ out,
    int N, int D, int H, int W)
{
    long long total = (long long)N * D * H * W;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int pw = idx % W;
    int ph = (idx / W) % H;
    int pd = (idx / (long long)(W * H)) % D;
    int n  = idx / (long long)(W * H * D);

    float vals[16];

    #pragma unroll
    for (int co = 0; co < 16; ++co) {
        float v0 = bias[co];
        float v1 = bias[co];
        float v2 = bias[co];
        float v3 = bias[co];
        float v4 = bias[co];
        float v5 = bias[co];
        float v6 = bias[co];
        float v7 = bias[co];

        #pragma unroll
        for (int ci = 0; ci < 3; ++ci) {
            const float* xp000 = x + (((long long)n * 3 + ci) * D + pd) * H * W + ph * W + pw;
            const float x000 = *xp000;
            const long long wbase = (((long long)ci * 16 + co) * 3) * 3 * 3;
            const float w111 = w[wbase + 13];
            const float w112 = w[wbase + 14];
            const float w110 = w[wbase + 12];
            const float w121 = w[wbase + 16];
            const float w101 = w[wbase + 10];
            const float w122 = w[wbase + 17];
            const float w120 = w[wbase + 15];
            const float w102 = w[wbase + 11];
            const float w100 = w[wbase + 9];
            const float w211 = w[wbase + 22];
            const float w011 = w[wbase + 4];
            const float w212 = w[wbase + 23];
            const float w210 = w[wbase + 21];
            const float w012 = w[wbase + 5];
            const float w010 = w[wbase + 3];
            const float w221 = w[wbase + 25];
            const float w201 = w[wbase + 19];
            const float w021 = w[wbase + 7];
            const float w001 = w[wbase + 1];
            const float w222 = w[wbase + 26];
            const float w220 = w[wbase + 24];
            const float w202 = w[wbase + 20];
            const float w200 = w[wbase + 18];
            const float w022 = w[wbase + 8];
            const float w020 = w[wbase + 6];
            const float w002 = w[wbase + 2];
            const float w000 = w[wbase + 0];

            v0 = fmaf(x000, w111, v0);
            v1 = fmaf(x000, w112, v1);
            v2 = fmaf(x000, w121, v2);
            v3 = fmaf(x000, w122, v3);
            v4 = fmaf(x000, w211, v4);
            v5 = fmaf(x000, w212, v5);
            v6 = fmaf(x000, w221, v6);
            v7 = fmaf(x000, w222, v7);

            if (pw + 1 < W) {
                float x001 = xp000[1];
                v1 = fmaf(x001, w110, v1);
                v3 = fmaf(x001, w120, v3);
                v5 = fmaf(x001, w210, v5);
                v7 = fmaf(x001, w220, v7);
            }
            if (ph + 1 < H) {
                float x010 = xp000[W];
                v2 = fmaf(x010, w101, v2);
                v3 = fmaf(x010, w102, v3);
                v6 = fmaf(x010, w201, v6);
                v7 = fmaf(x010, w202, v7);
            }
            if (pd + 1 < D) {
                float x100 = xp000[(long long)H*W];
                v4 = fmaf(x100, w011, v4);
                v5 = fmaf(x100, w012, v5);
                v6 = fmaf(x100, w021, v6);
                v7 = fmaf(x100, w022, v7);
            }
            if (ph + 1 < H && pw + 1 < W) {
                float x011 = xp000[W + 1];
                v3 = fmaf(x011, w100, v3);
                v7 = fmaf(x011, w200, v7);
            }
            if (pd + 1 < D && pw + 1 < W) {
                float x101 = xp000[(long long)H*W + 1];
                v5 = fmaf(x101, w010, v5);
                v7 = fmaf(x101, w020, v7);
            }
            if (pd + 1 < D && ph + 1 < H) {
                float x110 = xp000[(long long)H*W + W];
                v6 = fmaf(x110, w001, v6);
                v7 = fmaf(x110, w002, v7);
            }
            if (pd + 1 < D && ph + 1 < H && pw + 1 < W) {
                float x111x = xp000[(long long)H*W + W + 1];
                v7 = fmaf(x111x, w000, v7);
            }
        }
        float m1 = fmaxf(fmaxf(v0, v1), fmaxf(v2, v3));
        float m2 = fmaxf(fmaxf(v4, v5), fmaxf(v6, v7));
        vals[co] = fmaxf(m1, m2);
    }

    float m = vals[0];
    #pragma unroll
    for (int co = 1; co < 16; ++co) m = fmaxf(m, vals[co]);
    float sum = 0.0f;
    #pragma unroll
    for (int co = 0; co < 16; ++co) {
        float e = expf(vals[co] - m);
        vals[co] = e;
        sum += e;
    }
    float invsum = 1.0f / sum;
    float best = -FLT_MAX;
    #pragma unroll
    for (int co = 0; co < 16; ++co) {
        float z = vals[co] * invsum - sub[co];
        float sw = z / (1.0f + expf(-z));
        best = fmaxf(best, sw);
    }
    out[idx] = best;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub) {
    int N = x.size(0);
    int D = x.size(2);
    int H = x.size(3);
    int W = x.size(4);
    auto out = torch::empty({N, D, H, W}, x.options());
    long long total = (long long)N * D * H * W;
    int threads = 128;
    int blocks = (int)((total + threads - 1) / threads);
    fused_deconv_pool_softmax_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_softmax_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['fused_forward'], with_cuda=True, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_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(ModelNew, 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))
    def forward(self, x):
        return _ext.fused_forward(x.contiguous(), self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous(), self.subtract.contiguous())