```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Fully fused TF32 tensor-core kernel (v5, fixed):
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Fixes vs. the previous (crashing / incorrect) version:
#  1. The shared-memory parity tables are gone. The (ci, tap, offset) mapping is
#     now computed with integer arithmetic inside fully unrolled loops, so every
#     index folds to a compile-time immediate. The old table init wrote p=idx>>5
#     (32 per parity) into [8][24] arrays, leaving parities 6..7 uninitialized
#     garbage -> illegal global loads. That was the crash.
#  2. The epilogue no longer mixes the two output columns held by each thread.
#     Softmax denominators and the final channel-max are now reduced separately
#     for column 2*tg and column 2*tg+1 (the old code summed both columns into
#     one softmax denominator, which is mathematically wrong).
#  3. TF32 operands are rounded with cvt.rna.tf32.f32 (same conversion cuBLAS
#     uses) before mma.sync, keeping accuracy in the allowed TF32 class.
#
# Design:
#  * The stride-2 conv-transpose decomposes into 8 stride-1 "parity"
#    sub-convolutions on the input grid. The 2x2x2 pool window corresponds
#    exactly to the 8 parity classes at one grid cell, so pooling is an
#    elementwise max across the 8 accumulator fragments and the pooled spatial
#    size equals the input spatial size.
#  * Each warp computes a 16(co) x 8(cells) tile per parity with
#    mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32. All 13 weight fragments
#    are preloaded into registers once per thread and reused for every tile.
#  * B-fragments (gathered inputs) load straight from global memory with
#    per-thread addresses (offsets are compile-time constants after unrolling).
#  * Softmax exponentials are reused for the probabilities; subtract + swish +
#    channel-max run in registers with __shfl_xor butterflies over 8-lane
#    groups. No intermediate tensor is ever materialized in global memory.

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>

#define CI 3
#define CO 16
#define NFRAG 13

__device__ __forceinline__ unsigned f2tf32(const float f) {
    unsigned u;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(u) : "f"(f));
    return u;
}

struct TapInfo { int ci; int tap; int od; int oh; int ow; };

// Map the k-th reduction index of parity class p to (input channel, kernel tap,
// input offset). With p and k compile-time constants (unrolled loops) this
// folds to immediates.
__device__ __forceinline__ TapInfo tap_info(const int p, const int k) {
    const int kd = (p >> 2) & 1, kh = (p >> 1) & 1, kw = p & 1;
    const int TW = kw + 1, TH = kh + 1, TD = kd + 1;
    const int TPS = TD * TH * TW;
    TapInfo t;
    t.ci = 0; t.tap = 0; t.od = 0; t.oh = 0; t.ow = 0;
    if (k < CI * TPS) {
        t.ci = k / TPS;
        int r = k - t.ci * TPS;
        const int di = r / (TH * TW);
        r -= di * (TH * TW);
        const int hi = r / TW;
        const int wi = r - hi * TW;
        t.tap = (kd ? (di ? 2 : 0) : 1) * 9
              + (kh ? (hi ? 2 : 0) : 1) * 3
              + (kw ? (wi ? 2 : 0) : 1);
        t.od = kd ? (1 - di) : 0;
        t.oh = kh ? (1 - hi) : 0;
        t.ow = kw ? (1 - wi) : 0;
    }
    return t;
}

__device__ __forceinline__ int ktaps(const int p) {
    const int kd = (p >> 2) & 1, kh = (p >> 1) & 1, kw = p & 1;
    return CI * (kd + 1) * (kh + 1) * (kw + 1);
}

__global__ __launch_bounds__(128) void fused_convT_pool_softmax_swish_cmax_mma_kernel(
    const float* __restrict__ x,        // [B, CI, D, H, W]
    const float* __restrict__ weight,   // [CI, CO, 3, 3, 3]
    const float* __restrict__ bias,     // [CO]
    const float* __restrict__ sub,      // [CO]
    float* __restrict__ out,            // [B, D, H, W]  (pooled size == input size)
    const int D, const int H, const int W,
    const int nwt,                      // ceil(W/8) tiles per (d,h)
    const int tiles_per_b,              // D*H*nwt
    const int nblocks, const int T)     // blocks per batch, tiles per warp chunk
{
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    const int g  = lane >> 2;    // mma group id (0..7): rows g and g+8
    const int tg = lane & 3;     // thread id in group (0..3)

    const int HW  = H * W;
    const int DHW = D * HW;

    // ---- Preload all 13 A (weight) fragments into registers, TF32-rounded. ----
    unsigned a[NFRAG][4];
    {
        int f = 0;
        #pragma unroll
        for (int p = 0; p < 8; ++p) {
            const int Kp = ktaps(p);
            const int steps = (Kp + 7) >> 3;
            #pragma unroll
            for (int st = 0; st < 3; ++st) {
                if (st < steps) {
                    const int ka = 8 * st + tg;
                    const int kb = ka + 4;
                    const TapInfo ta = tap_info(p, ka);
                    const TapInfo tb = tap_info(p, kb);
                    const float* wg = weight + g * 27;
                    a[f][0] = (ka < Kp) ? f2tf32(__ldg(wg + ta.ci * 432 + ta.tap)) : 0u;
                    a[f][1] = (ka < Kp) ? f2tf32(__ldg(wg + 216 + ta.ci * 432 + ta.tap)) : 0u;
                    a[f][2] = (kb < Kp) ? f2tf32(__ldg(wg + tb.ci * 432 + tb.tap)) : 0u;
                    a[f][3] = (kb < Kp) ? f2tf32(__ldg(wg + 216 + tb.ci * 432 + tb.tap)) : 0u;
                    ++f;
                }
            }
        }
    }

    const float b0c = __ldg(bias + g);
    const float b8c = __ldg(bias + g + 8);
    const float s0c = __ldg(sub + g);
    const float s8c = __ldg(sub + g + 8);

    const int b = blockIdx.z;
    const float* xb = x + (long long)b * CI * DHW;
    float* outb = out + (long long)b * DHW;

    // ---- Tile loop: each warp handles a contiguous chunk of T tiles. ----
    for (int blk = blockIdx.x; blk < nblocks; blk += gridDim.x) {
        const int base = (blk * 4 + wid) * T;
        if (base >= tiles_per_b) continue;
        const int rem = tiles_per_b - base;
        const int tcount = rem < T ? rem : T;

        for (int t = 0; t < tcount; ++t) {
            const int tile = base + t;
            const int wt = tile % nwt;
            const int dh = tile / nwt;
            const int h = dh % H;
            const int d = dh / H;
            const int w0 = wt * 8;

            // Accumulators: [parity][c0..c3]
            float acc[8][4];
            #pragma unroll
            for (int p = 0; p < 8; ++p) {
                acc[p][0] = 0.0f; acc[p][1] = 0.0f;
                acc[p][2] = 0.0f; acc[p][3] = 0.0f;
            }

            int f = 0;
            #pragma unroll
            for (int p = 0; p < 8; ++p) {
                const int Kp = ktaps(p);
                const int steps = (Kp + 7) >> 3;
                #pragma unroll
                for (int st = 0; st < 3; ++st) {
                    if (st < steps) {
                        const int ka = 8 * st + tg;
                        const int kb = ka + 4;
                        const TapInfo ta = tap_info(p, ka);
                        const TapInfo tb = tap_info(p, kb);
                        // B fragment: row k, column g -> input w index w0 + g + offw
                        const int wa = w0 + g + ta.ow;
                        const int wb = w0 + g + tb.ow;
                        const int da = d + ta.od, ha = h + ta.oh;
                        const int db = d + tb.od, hb = h + tb.oh;
                        const bool va = (ka < Kp) && ((unsigned)wa < (unsigned)W)
                                      && ((unsigned)da < (unsigned)D)
                                      && ((unsigned)ha < (unsigned)H);
                        const bool vb = (kb < Kp) && ((unsigned)wb < (unsigned)W)
                                      && ((unsigned)db < (unsigned)D)
                                      && ((unsigned)hb < (unsigned)H);
                        const unsigned bv0 = va ? f2tf32(__ldg(xb + ta.ci * DHW
                                                     + da * HW + ha * W + wa)) : 0u;
                        const unsigned bv1 = vb ? f2tf32(__ldg(xb + tb.ci * DHW
                                                     + db * HW + hb * W + wb)) : 0u;
                        asm volatile(
                            "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
                            "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                            : "+f"(acc[p][0]), "+f"(acc[p][1]),
                              "+f"(acc[p][2]), "+f"(acc[p][3])
                            : "r"(a[f][0]), "r"(a[f][1]), "r"(a[f][2]), "r"(a[f][3]),
                              "r"(bv0), "r"(bv1));
                        ++f;
                    }
                }
            }

            // MaxPool3d across the 8 parity classes, then add the conv bias.
            // r0: (ch g, col 2tg)  r1: (ch g, col 2tg+1)
            // r2: (ch g+8, col 2tg)  r3: (ch g+8, col 2tg+1)
            float r0 = acc[0][0], r1 = acc[0][1], r2 = acc[0][2], r3 = acc[0][3];
            #pragma unroll
            for (int p = 1; p < 8; ++p) {
                r0 = fmaxf(r0, acc[p][0]);
                r1 = fmaxf(r1, acc[p][1]);
                r2 = fmaxf(r2, acc[p][2]);
                r3 = fmaxf(r3, acc[p][3]);
            }
            r0 += b0c; r1 += b0c; r2 += b8c; r3 += b8c;

            // Softmax across the 16 channels of each column (columns reduced
            // separately over the 8-lane groups with xor masks 4/8/16).
            float mA = fmaxf(r0, r2);
            float mB = fmaxf(r1, r3);
            mA = fmaxf(mA, __shfl_xor_sync(0xffffffffu, mA, 4));
            mA = fmaxf(mA, __shfl_xor_sync(0xffffffffu, mA, 8));
            mA = fmaxf(mA, __shfl_xor_sync(0xffffffffu, mA, 16));
            mB = fmaxf(mB, __shfl_xor_sync(0xffffffffu, mB, 4));
            mB = fmaxf(mB, __shfl_xor_sync(0xffffffffu, mB, 8));
            mB = fmaxf(mB, __shfl_xor_sync(0xffffffffu, mB, 16));

            const float eA0 = __expf(r0 - mA), eA1 = __expf(r2 - mA);
            const float eB0 = __expf(r1 - mB), eB1 = __expf(r3 - mB);
            float sA = eA0 + eA1, sB = eB0 + eB1;
            sA += __shfl_xor_sync(0xffffffffu, sA, 4);
            sA += __shfl_xor_sync(0xffffffffu, sA, 8);
            sA += __shfl_xor_sync(0xffffffffu, sA, 16);
            sB += __shfl_xor_sync(0xffffffffu, sB, 4);
            sB += __shfl_xor_sync(0xffffffffu, sB, 8);
            sB += __shfl_xor_sync(0xffffffffu, sB, 16);
            const float invA = __fdividef(1.0f, sA);
            const float invB = __fdividef(1.0f, sB);

            // Subtract per-channel constant, swish (y * sigmoid(y)), max over ch.
            const float yA0 = fmaf(eA0, invA, -s0c);
            const float yA1 = fmaf(eA1, invA, -s8c);
            const float yB0 = fmaf(eB0, invB, -s0c);
            const float yB1 = fmaf(eB1, invB, -s8c);
            const float zA0 = __fdividef(yA0, 1.0f + __expf(-yA0));
            const float zA1 = __fdividef(yA1, 1.0f + __expf(-yA1));
            const float zB0 = __fdividef(yB0, 1.0f + __expf(-yB0));
            const float zB1 = __fdividef(yB1, 1.0f + __expf(-yB1));
            float zA = fmaxf(zA0, zA1);
            float zB = fmaxf(zB0, zB1);
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 4));
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 8));
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 16));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 4));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 8));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 16));

            if (g == 0) {
                const int wo = w0 + 2 * tg;
                if (wo < W)     outb[d * HW + h * W + wo]     = zA;
                if (wo + 1 < W) outb[d * HW + h * W + wo + 1] = zB;
            }
        }
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            torch::Tensor sub)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() && sub.is_cuda(),
                "all tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.size(0) == CI && weight.size(1) == CO && weight.size(2) == 3,
                "unexpected weight shape");

    const int64_t B = x.size(0);
    const int64_t D = x.size(2);
    const int64_t H = x.size(3);
    const int64_t W = x.size(4);

    auto out = at::empty({B, D, H, W}, x.options());
    if (out.numel() == 0) return out;

    const int nwt = (int)((W + 7) >> 3);
    const int tiles_per_b = (int)(D * H) * nwt;

    static int nsm = 0;
    if (nsm == 0) {
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, dev);
    }
    int T = (tiles_per_b + nsm * 16 - 1) / (nsm * 16);
    if (T < 1) T = 1;
    const int nblocks = (tiles_per_b + 4 * T - 1) / (4 * T);
    const unsigned gx = nblocks > 65535 ? 65535u : (unsigned)nblocks;

    const dim3 grid(gx, 1u, (unsigned)B);
    const dim3 block(128, 1, 1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_convT_pool_softmax_swish_cmax_mma_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        sub.data_ptr<float>(), out.data_ptr<float>(),
        (int)D, (int)H, (int)W, nwt, tiles_per_b, nblocks, T);
    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, torch::Tensor sub);"
)

_fused_ext = load_inline(
    name="fused_convT3d_pool_softmax_sub_swish_cmax_mma_v5",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Uses a single fully-fused TF32 tensor-core CUDA
    kernel that performs ConvTranspose3d + MaxPool3d + Softmax + Subtract +
    Swish + channel-wise Max without materializing any intermediate tensor in
    global memory. Falls back to the original PyTorch ops for unsupported
    configurations.
    """
    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 _i(v):
            return v[0] if isinstance(v, (tuple, list)) else v

        cfg = (in_channels, out_channels, _i(kernel_size), _i(stride), _i(padding),
               _i(output_padding), _i(pool_kernel_size), _i(pool_stride), _i(pool_padding))
        self._fused_ok = (cfg == (3, 16, 3, 2, 1, 1, 2, 2, 0))
        self._cap_ok = None
        self._zero_bias = None

    def forward(self, x):
        if self._fused_ok and x.is_cuda and x.dtype == torch.float32:
            if self._cap_ok is None:
                self._cap_ok = torch.cuda.get_device_capability(x.device)[0] >= 8
            if self._cap_ok:
                B, D, H, W = x.shape[0], x.shape[2], x.shape[3], x.shape[4]
                # conv-transpose output size: (n-1)*2 - 2 + 3 + 1 = 2n
                Dc = (D - 1) * 2 - 2 + 3 + 1
                Hc = (H - 1) * 2 - 2 + 3 + 1
                Wc = (W - 1) * 2 - 2 + 3 + 1
                # max-pool (k=2, s=2, p=0) output size
                Dp = (Dc - 2) // 2 + 1
                Hp = (Hc - 2) // 2 + 1
                Wp = (Wc - 2) // 2 + 1
                # parity decomposition requires pooled size == input size
                if (Dp == D and Hp == H and Wp == W and B <= 65535
                        and x.numel() < 2**31 and B * D * H * W < 2**31):
                    bias = self.conv_transpose.bias
                    if bias is None:
                        if self._zero_bias is None or self._zero_bias.device != x.device:
                            self._zero_bias = torch.zeros(16, device=x.device,
                                                          dtype=torch.float32)
                        bias = self._zero_bias
                    return _fused_ext.fused_forward(
                        x.contiguous(),
                        self.conv_transpose.weight.contiguous(),
                        bias.contiguous(),
                        self.subtract.contiguous())

        # Generic fallback: original PyTorch operator sequence.
        x = self.conv_transpose(x)
        x = self.max_pool(x)
        x = torch.softmax(x, dim=1)
        x = x - self.subtract.view(1, -1, 1, 1, 1)
        x = torch.sigmoid(x) * x
        x = torch.max(x, dim=1)[0]
        return x
```