import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies Softmax and Sigmoid.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=True):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, bias=bias)
        self.softmax = nn.Softmax(dim=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, D, H, W).
        """
        x = self.conv_transpose(x)
        x = self.softmax(x)
        x = self.sigmoid(x)
        return x

batch_size = 16
in_channels = 32
out_channels = 64
D, H, W = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_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 convt3d_softmax_sigmoid_forward(torch::Tensor x,
                                               torch::Tensor w,
                                               torch::Tensor b,
                                               int stride,
                                               int padding,
                                               int output_padding);
'''

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

__global__ void convt3d_s2p1_k3_softmax_sigmoid_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ out,
    int N, int Ci, int ID, int IH, int IW,
    int Co, int OD, int OH, int OW,
    bool has_bias)
{
    // Specialised for the benchmark shape: Co=64, k=3, stride=2, padding=1.
    // Four independent output spatial locations per CTA; each 64-thread group
    // computes all channels for one location, then performs the channel softmax.
    int tid = threadIdx.x;          // 0..255
    int group = tid >> 6;           // 0..3
    int co = tid & 63;              // channel within group
    int warp_lane = tid & 31;
    int warp_in_group = (co >> 5);  // 0 or 1

    int total_pos = N * OD * OH * OW;
    int pos = blockIdx.x * 4 + group;
    bool valid = (pos < total_pos) && (co < Co);

    int n = 0, od = 0, oh = 0, ow = 0;
    if (pos < total_pos) {
        int t = pos;
        ow = t % OW; t /= OW;
        oh = t % OH; t /= OH;
        od = t % OD; t /= OD;
        n = t;
    }

    float acc = valid ? (has_bias ? __ldg(b + co) : 0.0f) : -FLT_MAX;

    if (valid) {
        int didx[2], dk[2], hidx[2], hk[2], widx[2], wk[2];
        int nd = 0, nh = 0, nw = 0;

        if ((od & 1) == 0) {
            int id = od >> 1;
            if (id < ID) { didx[0] = id; dk[0] = 1; nd = 1; }
        } else {
            int id0 = (od + 1) >> 1;
            if (id0 < ID) { didx[nd] = id0; dk[nd] = 0; nd++; }
            int id1 = (od - 1) >> 1;
            if (id1 >= 0) { didx[nd] = id1; dk[nd] = 2; nd++; }
        }

        if ((oh & 1) == 0) {
            int ih = oh >> 1;
            if (ih < IH) { hidx[0] = ih; hk[0] = 1; nh = 1; }
        } else {
            int ih0 = (oh + 1) >> 1;
            if (ih0 < IH) { hidx[nh] = ih0; hk[nh] = 0; nh++; }
            int ih1 = (oh - 1) >> 1;
            if (ih1 >= 0) { hidx[nh] = ih1; hk[nh] = 2; nh++; }
        }

        if ((ow & 1) == 0) {
            int iw = ow >> 1;
            if (iw < IW) { widx[0] = iw; wk[0] = 1; nw = 1; }
        } else {
            int iw0 = (ow + 1) >> 1;
            if (iw0 < IW) { widx[nw] = iw0; wk[nw] = 0; nw++; }
            int iw1 = (ow - 1) >> 1;
            if (iw1 >= 0) { widx[nw] = iw1; wk[nw] = 2; nw++; }
        }

        int n_base_x = n * Ci * ID * IH * IW;
        #pragma unroll 1
        for (int ci = 0; ci < Ci; ++ci) {
            int ci_x_base = n_base_x + ci * ID * IH * IW;
            int ci_w_base = ci * Co * 27 + co * 27;
            #pragma unroll
            for (int a = 0; a < 2; ++a) {
                if (a >= nd) continue;
                int ix_d_base = ci_x_base + didx[a] * IH * IW;
                int kd_off = dk[a] * 9;
                #pragma unroll
                for (int bb = 0; bb < 2; ++bb) {
                    if (bb >= nh) continue;
                    int ix_h_base = ix_d_base + hidx[bb] * IW;
                    int kh_off = hk[bb] * 3;
                    #pragma unroll
                    for (int c = 0; c < 2; ++c) {
                        if (c >= nw) continue;
                        float xv = __ldg(x + ix_h_base + widx[c]);
                        float wv = __ldg(w + ci_w_base + kd_off + kh_off + wk[c]);
                        acc = fmaf(xv, wv, acc);
                    }
                }
            }
        }
    }

    __shared__ float sh[8]; // 4 groups * 2 warps: reused for max and sum

    float m = acc;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        m = fmaxf(m, __shfl_down_sync(0xffffffff, m, off));
    if (warp_lane == 0) sh[group * 2 + warp_in_group] = m;
    __syncthreads();

    float maxv = fmaxf(sh[group * 2], sh[group * 2 + 1]);
    float e = valid ? __expf(acc - maxv) : 0.0f;

    float s = e;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(0xffffffff, s, off);
    if (warp_lane == 0) sh[group * 2 + warp_in_group] = s;
    __syncthreads();

    float denom = sh[group * 2] + sh[group * 2 + 1];
    if (valid) {
        float sm = e / denom;
        float res = 1.0f / (1.0f + __expf(-sm));
        int out_idx = (((n * Co + co) * OD + od) * OH + oh) * OW + ow;
        out[out_idx] = res;
    }
}

torch::Tensor convt3d_softmax_sigmoid_forward(torch::Tensor x,
                                               torch::Tensor w,
                                               torch::Tensor b,
                                               int stride,
                                               int padding,
                                               int output_padding)
{
    x = x.contiguous();
    w = w.contiguous();
    const int N  = (int)x.size(0);
    const int Ci = (int)x.size(1);
    const int ID = (int)x.size(2);
    const int IH = (int)x.size(3);
    const int IW = (int)x.size(4);
    const int Co = (int)w.size(1);
    const int KD = (int)w.size(2);
    const int KH = (int)w.size(3);
    const int KW = (int)w.size(4);
    const int OD = (ID - 1) * stride - 2 * padding + KD + output_padding;
    const int OH = (IH - 1) * stride - 2 * padding + KH + output_padding;
    const int OW = (IW - 1) * stride - 2 * padding + KW + output_padding;

    auto out = torch::empty({N, Co, OD, OH, OW}, x.options());
    int total_pos = N * OD * OH * OW;
    int blocks = (total_pos + 3) / 4;
    bool has_bias = (b.numel() != 0);

    convt3d_s2p1_k3_softmax_sigmoid_kernel<<<blocks, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? b.data_ptr<float>() : nullptr,
        out.data_ptr<float>(), N, Ci, ID, IH, IW, Co, OD, OH, OW, has_bias);
    return out;
}
'''

_ext = load_inline(
    name='convt3d_softmax_sigmoid_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['convt3d_softmax_sigmoid_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, output_padding=output_padding, bias=bias
        )
        self._stride = self.conv_transpose.stride[0]
        self._padding = self.conv_transpose.padding[0]
        self._output_padding = self.conv_transpose.output_padding[0]
        self.register_buffer('_empty_bias', torch.empty(0), persistent=False)

    def forward(self, x):
        b = self.conv_transpose.bias
        if b is None:
            b = self._empty_bias.to(device=x.device, dtype=x.dtype)
        return _ext.convt3d_softmax_sigmoid_forward(
            x, self.conv_transpose.weight, b,
            int(self._stride), int(self._padding), int(self._output_padding)
        )