```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 (v6):
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Design notes:
#  * 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 a
#    per-parity max over accumulator fragments and the pooled spatial size
#    equals the input spatial size (16 x 32 x 32 here).
#  * mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 with M = 16 pooled
#    voxels (consecutive along W) and N = 8 channels (two mmas cover the 16
#    output channels). Putting channels on N makes the cross-channel softmax /
#    final max reduce over the two low lane bits (xor 1, 2) instead of the
#    three high bits, and pairs of voxels are reduced together, which cuts the
#    epilogue shuffles ~3x versus the M=channel layout.
#  * All 52 weight fragment registers are preloaded once per thread (TF32
#    rounded with cvt.rna.tf32.f32, the same rounding cuBLAS uses) and reused
#    for every tile. Accumulators are zeroed per parity and folded into the
#    running pool-max immediately, so only 8 C + 8 pooled registers are live.
#  * Input fragment loads use one base register per tile with compile-time
#    immediate offsets; bounds predicates (zero padding) reduce to three
#    hoisted per-tile booleans combined with compile-time tap flags.
#  * The epilogue reuses the softmax exponentials for the probabilities and
#    evaluates swish as h*(1 + tanh(h)), h = 0.5*y, with the hardware
#    tanh.approx.f32 instruction (error far inside the TF32 tolerance class).
#  * 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 NSTEP 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;
}

__device__ __forceinline__ float tanh_apx(const float x) {
    float y;
    asm volatile("tanh.approx.f32 %0, %1;\n" : "=f"(y) : "f"(x));
    return y;
}

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 (fully 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__(256, 2)
void fused_convT_pool_softmax_swish_cmax_mma_v6_kernel(
    const float* __restrict__ x,        // [B, 3, 16, 32, 32]
    const float* __restrict__ weight,   // [3, 16, 3, 3, 3]
    const float* __restrict__ bias,     // [16]
    const float* __restrict__ sub,      // [16]
    float* __restrict__ out,            // [B, 16, 32, 32]
    const int total_tiles,              // B * 1024
    const int T)                        // 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): voxel rows g, g+8
    const int tg = lane & 3;     // thread id in group (0..3)

    // ---- Preload all weight fragments (TF32 rounded) once per thread. ----
    // Frag f (parity p, k-step st) holds:
    //   [0] = w[ci(tg)][g  ][tap(tg)],   [1] = w[ci(tg+4)][g  ][tap(tg+4)]
    //   [2] = w[ci(tg)][g+8][tap(tg)],   [3] = w[ci(tg+4)][g+8][tap(tg+4)]
    unsigned wf[NSTEP][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);
                    wf[f][0] = (ka < Kp) ? f2tf32(__ldg(weight + ta.ci * 432 + g * 27       + ta.tap)) : 0u;
                    wf[f][1] = (kb < Kp) ? f2tf32(__ldg(weight + tb.ci * 432 + g * 27       + tb.tap)) : 0u;
                    wf[f][2] = (ka < Kp) ? f2tf32(__ldg(weight + ta.ci * 432 + (g + 8) * 27 + ta.tap)) : 0u;
                    wf[f][3] = (kb < Kp) ? f2tf32(__ldg(weight + tb.ci * 432 + (g + 8) * 27 + tb.tap)) : 0u;
                    ++f;
                }
            }
        }
    }

    // Per-thread constant bias / subtract values for its four channels.
    const float b0c = __ldg(bias + 2 * tg);
    const float b1c = __ldg(bias + 2 * tg + 1);
    const float b2c = __ldg(bias + 8 + 2 * tg);
    const float b3c = __ldg(bias + 8 + 2 * tg + 1);
    const float s0c = __ldg(sub + 2 * tg);
    const float s1c = __ldg(sub + 2 * tg + 1);
    const float s2c = __ldg(sub + 8 + 2 * tg);
    const float s3c = __ldg(sub + 8 + 2 * tg + 1);

    // ---- Chunked persistent loop over tiles (tile = 16 voxels along W). ----
    const int gwarp = blockIdx.x * 8 + wid;
    const int start = gwarp * T;
    const int end = (total_tiles - start < T) ? total_tiles : start + T;

    for (int work = start; work < end; ++work) {
        const int b    = work >> 10;          // 1024 tiles per batch element
        const int tile = work & 1023;
        const int wt   = tile & 1;            // 2 tiles of 16 along W
        const int dh   = tile >> 1;
        const int h    = dh & 31;
        const int d    = dh >> 5;
        const int w0   = wt << 4;

        const float* xp = x + b * 49152 + d * 1024 + h * 32 + (w0 + g);
        const bool dok = (d < 15);
        const bool hok = (h < 31);
        const bool wok = (w0 + g + 1 < 32);

        float r[8];
        #pragma unroll
        for (int j = 0; j < 8; ++j) r[j] = -FLT_MAX;

        int f = 0;
        #pragma unroll
        for (int p = 0; p < 8; ++p) {
            const int Kp = ktaps(p);
            const int steps = (Kp + 7) >> 3;
            float acc[8];
            #pragma unroll
            for (int j = 0; j < 8; ++j) acc[j] = 0.0f;
            #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 int AOFF = ta.ci * 16384 + ta.od * 1024 + ta.oh * 32 + ta.ow;
                    const int BOFF = tb.ci * 16384 + tb.od * 1024 + tb.oh * 32 + tb.ow;
                    const bool pa = (ka < Kp)
                        && (ta.od ? dok : true) && (ta.oh ? hok : true)
                        && (ta.ow ? wok : true);
                    const bool pb = (kb < Kp)
                        && (tb.od ? dok : true) && (tb.oh ? hok : true)
                        && (tb.ow ? wok : true);
                    unsigned x0 = 0u, x1 = 0u, x2 = 0u, x3 = 0u;
                    if (pa) { x0 = f2tf32(__ldg(xp + AOFF));
                              x1 = f2tf32(__ldg(xp + AOFF + 8)); }
                    if (pb) { x2 = f2tf32(__ldg(xp + BOFF));
                              x3 = f2tf32(__ldg(xp + BOFF + 8)); }
                    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[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3])
                        : "r"(x0), "r"(x1), "r"(x2), "r"(x3),
                          "r"(wf[f][0]), "r"(wf[f][1]));
                    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[4]), "+f"(acc[5]), "+f"(acc[6]), "+f"(acc[7])
                        : "r"(x0), "r"(x1), "r"(x2), "r"(x3),
                          "r"(wf[f][2]), "r"(wf[f][3]));
                    ++f;
                }
            }
            #pragma unroll
            for (int j = 0; j < 8; ++j) r[j] = fmaxf(r[j], acc[j]);
        }

        // Add the conv bias (folded after the pool max).
        r[0] += b0c; r[1] += b1c; r[2] += b0c; r[3] += b1c;
        r[4] += b2c; r[5] += b3c; r[6] += b2c; r[7] += b3c;

        // Softmax across the 16 channels of each voxel. Voxel A (row g) holds
        // channels {2tg, 2tg+1, 8+2tg, 8+2tg+1} in r[0],r[1],r[4],r[5];
        // voxel B (row g+8) holds them in r[2],r[3],r[6],r[7]. The remaining
        // channels live in lanes with the same g (xor masks 1 and 2).
        float mA = fmaxf(fmaxf(r[0], r[1]), fmaxf(r[4], r[5]));
        float mB = fmaxf(fmaxf(r[2], r[3]), fmaxf(r[6], r[7]));
        mA = fmaxf(mA, __shfl_xor_sync(0xffffffffu, mA, 1));
        mA = fmaxf(mA, __shfl_xor_sync(0xffffffffu, mA, 2));
        mB = fmaxf(mB, __shfl_xor_sync(0xffffffffu, mB, 1));
        mB = fmaxf(mB, __shfl_xor_sync(0xffffffffu, mB, 2));

        const float eA0 = __expf(r[0] - mA), eA1 = __expf(r[1] - mA);
        const float eA2 = __expf(r[4] - mA), eA3 = __expf(r[5] - mA);
        const float eB0 = __expf(r[2] - mB), eB1 = __expf(r[3] - mB);
        const float eB2 = __expf(r[6] - mB), eB3 = __expf(r[7] - mB);
        float sA = (eA0 + eA1) + (eA2 + eA3);
        float sB = (eB0 + eB1) + (eB2 + eB3);
        sA += __shfl_xor_sync(0xffffffffu, sA, 1);
        sA += __shfl_xor_sync(0xffffffffu, sA, 2);
        sB += __shfl_xor_sync(0xffffffffu, sB, 1);
        sB += __shfl_xor_sync(0xffffffffu, sB, 2);
        const float invA = __fdividef(1.0f, sA);
        const float invB = __fdividef(1.0f, sB);

        // Subtract per-channel constant, swish via tanh, max over channels.
        const float yA0 = fmaf(eA0, invA, -s0c);
        const float yA1 = fmaf(eA1, invA, -s1c);
        const float yA2 = fmaf(eA2, invA, -s2c);
        const float yA3 = fmaf(eA3, invA, -s3c);
        const float yB0 = fmaf(eB0, invB, -s0c);
        const float yB1 = fmaf(eB1, invB, -s1c);
        const float yB2 = fmaf(eB2, invB, -s2c);
        const float yB3 = fmaf(eB3, invB, -s3c);
        const float hA0 = 0.5f * yA0, hA1 = 0.5f * yA1;
        const float hA2 = 0.5f * yA2, hA3 = 0.5f * yA3;
        const float hB0 = 0.5f * yB0, hB1 = 0.5f * yB1;
        const float hB2 = 0.5f * yB2, hB3 = 0.5f * yB3;
        const float zA0 = fmaf(hA0, tanh_apx(hA0), hA0);
        const float zA1 = fmaf(hA1, tanh_apx(hA1), hA1);
        const float zA2 = fmaf(hA2, tanh_apx(hA2), hA2);
        const float zA3 = fmaf(hA3, tanh_apx(hA3), hA3);
        const float zB0 = fmaf(hB0, tanh_apx(hB0), hB0);
        const float zB1 = fmaf(hB1, tanh_apx(hB1), hB1);
        const float zB2 = fmaf(hB2, tanh_apx(hB2), hB2);
        const float zB3 = fmaf(hB3, tanh_apx(hB3), hB3);
        float zA = fmaxf(fmaxf(zA0, zA1), fmaxf(zA2, zA3));
        float zB = fmaxf(fmaxf(zB0, zB1), fmaxf(zB2, zB3));
        zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 1));
        zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 2));
        zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 1));
        zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 2));

        if (tg == 0) {
            float* op = out + b * 16384 + d * 1024 + h * 32 + w0 + g;
            op[0] = zA;
            op[8] = 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");

    const int64_t B = x.size(0);
    auto out = at::empty({B, 16, 32, 32}, x.options());
    if (out.numel() == 0) return out;

    const int total_tiles = (int)B * 1024;
    int T = (total_tiles + 2047) / 2048;      // ~2048 warps worth of chunks
    if (T < 8) T = 8;
    const int64_t blocks = ((int64_t)total_tiles + (int64_t)T * 8 - 1) / ((int64_t)T * 8);

    const dim3 grid((unsigned)blocks, 1u, 1u);
    const dim3 block(256, 1, 1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_convT_pool_softmax_swish_cmax_mma_v6_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>(), total_tiles, 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_v6",
    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]
                # The parity decomposition requires pooled size == input size;
                # for k=3, s=2, p=1, op=1 and pool (2,2,0) that means the
                # conv-transpose output is exactly 2n per dim. The kernel is
                # specialized to the benchmark geometry 16 x 32 x 32.
                Dc = (D - 1) * 2 - 2 + 3 + 1
                Hc = (H - 1) * 2 - 2 + 3 + 1
                Wc = (W - 1) * 2 - 2 + 3 + 1
                if (D == 16 and H == 32 and W == 32 and Dc == 2 * D
                        and Hc == 2 * H and Wc == 2 * W
                        and B * 1024 <= 2**30):
                    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
```