```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernel: ConvTranspose3d (stride=1, padding=0, output_padding=0,
# groups=1) implemented as an implicit GEMM on tensor cores (TF32 WMMA).
# The transposed convolution with stride 1 and no padding is mathematically
# equivalent to a correlation:
#   out[b,co,d,h,w] = sum_{ci,kd,kh,kw} x[b,ci,d-kd,h-kh,w-kw] * W[ci,co,kd,kh,kw]
# which is evaluated as C(M=CO, N=voxels) = A(CO, CI*K) * B(CI*K, voxels)
# with the im2col operand B gathered on the fly into shared memory.

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

using namespace nvcuda;

#define BM 64
#define BN 64
#define BK 16
#define NTHREADS 256
#define WARPS_N 4
#define WARP_M 32
#define WARP_N 16

__global__ void __launch_bounds__(NTHREADS)
convt3d_s1_kernel(const float* __restrict__ x,
                  const float* __restrict__ wt,
                  const float* __restrict__ bias,
                  float* __restrict__ out,
                  int B, int CI, int CO,
                  int ID, int IH, int IW,
                  int KD, int KH, int KW,
                  int OD, int OH, int OW,
                  int has_bias)
{
    const int KKK  = KD * KH * KW;
    const int K    = CI * KKK;
    const int HKW  = KH * KW;
    const int Nvox = OD * OH * OW;

    const int b  = blockIdx.z;
    const int m0 = blockIdx.y * BM;
    const int n0 = blockIdx.x * BN;

    __shared__ float As[BM * BK];
    __shared__ float Bs[BN * BK];
    __shared__ float Cs[BM * BN];
    __shared__ int4  tab[2][BK];

    const int tid     = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m  = warp_id / WARPS_N;
    const int warp_n  = warp_id % WARPS_N;

    // This thread's im2col column (voxel) -- fixed across the whole K loop.
    const int n_thread = tid % BN;
    const int vox_t    = n0 + n_thread;
    const bool vox_ok  = (vox_t < Nvox);
    int d_t = 0, h_t = 0, w_t = 0;
    if (vox_ok) {
        d_t = vox_t / (OH * OW);
        int r = vox_t - d_t * (OH * OW);
        h_t = r / OW;
        w_t = r - h_t * OW;
    }

    const float* xb = x + (long)b * CI * ID * IH * IW;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2];
    #pragma unroll
    for (int i = 0; i < 2; ++i) wmma::fill_fragment(acc[i], 0.0f);

    const int nchunks = (K + BK - 1) / BK;

    // Prologue: decomposition table (ci,kd,kh,kw) for k-chunk 0.
    if (tid < BK) {
        int kk = tid;
        int4 t;
        if (kk < K) {
            int ci  = kk / KKK;
            int tap = kk - ci * KKK;
            int kd  = tap / HKW;
            int rem = tap - kd * HKW;
            int kh  = rem / KW;
            int kw  = rem - kh * KW;
            t = make_int4(ci, kd, kh, kw);
        } else {
            t = make_int4(-1, 0, 0, 0);
        }
        tab[0][tid] = t;
    }
    __syncthreads();

    int cur = 0, nxt = 1;

    for (int c = 0; c < nchunks; ++c) {
        // ---- Stage A tile (BM x BK): weights ----
        #pragma unroll
        for (int i = 0; i < (BM * BK) / NTHREADS; ++i) {
            int idx = tid + i * NTHREADS;
            int m   = idx / BK;
            int r   = idx % BK;
            int4 t  = tab[cur][r];
            float v = 0.0f;
            int co  = m0 + m;
            if (t.x >= 0 && co < CO) {
                int tap = t.y * HKW + t.z * KW + t.w;
                v = wt[((long)t.x * CO + co) * KKK + tap];
            }
            As[m * BK + r] = v;
        }
        // ---- Stage B tile (BN x BK): gathered input (im2col on the fly) ----
        #pragma unroll
        for (int i = 0; i < (BN * BK) / NTHREADS; ++i) {
            int idx = tid + i * NTHREADS;
            int r   = idx / BN;   // k row
            int n   = idx % BN;   // voxel column (== n_thread)
            int4 t  = tab[cur][r];
            float v = 0.0f;
            if (t.x >= 0 && vox_ok) {
                int di = d_t - t.y;
                int hi = h_t - t.z;
                int wi = w_t - t.w;
                if (di >= 0 && di < ID && hi >= 0 && hi < IH && wi >= 0 && wi < IW) {
                    v = xb[(((long)t.x * ID + di) * IH + hi) * IW + wi];
                }
            }
            Bs[n * BK + r] = v;
        }
        __syncthreads();

        // ---- Tensor-core MMA over the k-chunk ----
        {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> b_frag;
            #pragma unroll
            for (int kk = 0; kk < BK; kk += 8) {
                wmma::load_matrix_sync(b_frag, &Bs[(warp_n * WARP_N) * BK + kk], BK);
                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    wmma::load_matrix_sync(a_frag[i], &As[(warp_m * WARP_M + i * 16) * BK + kk], BK);
                    wmma::mma_sync(acc[i], a_frag[i], b_frag, acc[i]);
                }
            }
        }

        // ---- Decomposition table for the next chunk (double buffered) ----
        if (c + 1 < nchunks && tid < BK) {
            int kk = (c + 1) * BK + tid;
            int4 t;
            if (kk < K) {
                int ci  = kk / KKK;
                int tap = kk - ci * KKK;
                int kd  = tap / HKW;
                int rem = tap - kd * HKW;
                int kh  = rem / KW;
                int kw  = rem - kh * KW;
                t = make_int4(ci, kd, kh, kw);
            } else {
                t = make_int4(-1, 0, 0, 0);
            }
            tab[nxt][tid] = t;
        }
        __syncthreads();
        cur ^= 1;
        nxt ^= 1;
    }

    // ---- Epilogue: write the 64x64 output tile ----
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        wmma::store_matrix_sync(&Cs[(warp_m * WARP_M + i * 16) * BN + warp_n * WARP_N],
                                acc[i], BN, wmma::mem_row_major);
    }
    __syncthreads();

    float* outb = out + (long)b * CO * Nvox;
    #pragma unroll
    for (int i = 0; i < (BM * BN) / NTHREADS; ++i) {
        int idx = tid + i * NTHREADS;
        int m   = idx / BN;
        int n   = idx % BN;
        int co  = m0 + m;
        int vox = n0 + n;
        if (co < CO && vox < Nvox) {
            float val = Cs[m * BN + n];
            if (has_bias) val += bias[co];
            outb[(long)co * Nvox + vox] = val;
        }
    }
}

torch::Tensor convt3d_forward(torch::Tensor x, torch::Tensor weight,
                              torch::Tensor bias, bool has_bias)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    const int B  = xc.size(0);
    const int CI = xc.size(1);
    const int ID = xc.size(2);
    const int IH = xc.size(3);
    const int IW = xc.size(4);

    const int CO = wc.size(1);
    const int KD = wc.size(2);
    const int KH = wc.size(3);
    const int KW = wc.size(4);

    const int OD = ID + KD - 1;
    const int OH = IH + KH - 1;
    const int OW = IW + KW - 1;

    auto out = torch::empty({B, CO, OD, OH, OW}, xc.options());

    const int Nvox = OD * OH * OW;
    dim3 grid((Nvox + BN - 1) / BN, (CO + BM - 1) / BM, B);
    dim3 block(NTHREADS);

    const float* bias_ptr = has_bias ? bias.data_ptr<float>() : nullptr;

    convt3d_s1_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(),
        B, CI, CO, ID, IH, IW, KD, KH, KW, OD, OH, OW,
        has_bias ? 1 : 0);

    return out;
}
"""

_convt3d_cpp_src = (
    "torch::Tensor convt3d_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, bool has_bias);"
)

_convt3d_ext = load_inline(
    name="convt3d_tensorcore_ext",
    cpp_sources=_convt3d_cpp_src,
    cuda_sources=_convt3d_cuda_src,
    functions=["convt3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized transposed 3D convolution.

    Uses a custom tensor-core (TF32 WMMA) implicit-GEMM CUDA kernel when the
    configuration is supported (stride=1, padding=0, output_padding=0,
    groups=1, float32 CUDA input, compute capability >= 8.0) and falls back to
    the standard PyTorch operator otherwise.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0),
                 output_padding: tuple = (0, 0, 0), groups: int = 1,
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding,
            groups=groups, bias=bias)
        self._fast_path = (
            tuple(stride) == (1, 1, 1)
            and tuple(padding) == (0, 0, 0)
            and tuple(output_padding) == (0, 0, 0)
            and groups == 1
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (self._fast_path and x.is_cuda and x.dtype == torch.float32):
            cap = torch.cuda.get_device_capability(x.device)
            if cap[0] >= 8:
                conv = self.conv_transpose3d
                has_bias = conv.bias is not None
                bias = conv.bias if has_bias else torch.empty(
                    0, device=x.device, dtype=torch.float32)
                return _convt3d_ext.convt3d_forward(
                    x, conv.weight, bias, has_bias)
        return self.conv_transpose3d(x)
```