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

# Let cuDNN autotune the fastest engine for the fixed shapes/layouts.
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True

# v6: Fully fused direct conv3d (3x3x3, Cin=8, Cout=64, stride=1, pad=0)
# implemented with TF32 tensor-core MMA (mma.m16n8k8) and the entire epilogue
# (bias + LeakyReLU(0.2) + per-channel sum + clamp(-1,1) + exact GELU) fused
# into the accumulator store. It reads NCDHW input and writes NCDHW output
# directly: ONE kernel replaces layout converts + cuDNN conv + epilogue,
# removing ~3.5GB of intermediate global-memory traffic.
#   * Block: 128 threads (4 warps), tile = 64 cout x 64 voxels of one output
#     row (one do, one ho, 64 consecutive w).
#   * Weights are pre-packed on host into the exact m16n8k8 A-fragment layout
#     -> one aligned float4 load per thread per tap.
#   * The 9x8x66-float input halo is staged once in shared memory with a
#     bank-conflict-free pitch of 72.
#   * A generic cuDNN-based fallback path (v4 kernels) covers all other
#     configurations with identical numerics.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define NEG_SLOPE 0.2f
#define RSQRT2    0.70710678118654752440f  // 1 / sqrt(2)

__device__ __forceinline__ float epilogue_fn(float v, float b, float s) {
    v += b;                                        // conv bias
    v = v > 0.0f ? v : NEG_SLOPE * v;              // leaky_relu(0.2)
    v += s;                                        // broadcast per-channel sum
    v = fminf(fmaxf(v, -1.0f), 1.0f);              // clamp(-1, 1)
    return 0.5f * v * (1.0f + erff(v * RSQRT2));   // exact GELU
}

// ===========================================================================
// Fused direct conv3d: y = gelu(clamp(leaky_relu(conv3d(x,w)+bias)+sum,-1,1))
//   x:    NCDHW [N, 8, Di, Hi, Wi]   (contiguous, Wi % 4 == 0, Wo = Wi-2 <= 64)
//   wfrag4: packed A fragments [27][4][32] float4 (see host packing)
//   out:  NCDHW [N, 64, Do, Ho, Wo]  (contiguous)
// Block tile: 64 cout (all) x 64 voxels of one (do, ho) output row.
// Warp: 16 cout x 64 voxels = 8 m16n8k8 MMA n-tiles; k = 8 cin per tap.
// ===========================================================================
__global__ void __launch_bounds__(128)
conv3d_cin8_cout64_tf32_fused_kernel(const float* __restrict__ x,
                                     const float4* __restrict__ wfrag4,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ sumt,
                                     float* __restrict__ out,
                                     int Do, int Ho, int Wo,
                                     int Di, int Hi, int Wi) {
    // smem halo: sm[dz*3+dy][cin][w - w0], pitch 72 (conflict-free).
    __shared__ __align__(16) float sm[9 * 8 * 72];

    const int n = blockIdx.z / Do;
    const int do_ = blockIdx.z - n * Do;
    const int ho = blockIdx.y;
    const int w0 = blockIdx.x << 6;
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    const long HW = (long)Hi * Wi;
    const long DHW = (long)Di * HW;
    const long HWo = (long)Ho * Wo;
    const long DHWo = (long)Do * HWo;
    const long CPlane = DHWo;
    const long rowOff = (long)do_ * HWo + (long)ho * Wo;  // output row offset
    const long nInOff = (long)n * 8 * DHW;
    const long nOutOff = (long)n * 64 * CPlane;

    // ---- stage the input halo into smem (float4, guarded) ----
    {
        const int r = tid >> 4;        // 0..7
        const int c = tid & 15;        // float4 column
        #pragma unroll
        for (int it = 0; it < 9; ++it) {
            const int rr = r + (it << 3);   // 0..71
            const int cin = rr >> 3;        // 0..7
            const int rr8 = rr & 7;         // 0..7
            const int dz = rr8 / 3;
            const int dy = rr8 - dz * 3;
            const int w = w0 + (c << 2);
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (w + 3 < Wi) {
                const float* g = x + (((nInOff + (long)cin * DHW) +
                                       (long)(do_ + dz) * HW) +
                                       (long)(ho + dy) * Wi) + w;
                v = *reinterpret_cast<const float4*>(g);
            } else if (w < Wi) {
                const float* g = x + (((nInOff + (long)cin * DHW) +
                                       (long)(do_ + dz) * HW) +
                                       (long)(ho + dy) * Wi) + w;
                float tmp[4] = {0.f, 0.f, 0.f, 0.f};
                for (int j = 0; j < Wi - w; ++j) tmp[j] = g[j];
                v = make_float4(tmp[0], tmp[1], tmp[2], tmp[3]);
            }
            *reinterpret_cast<float4*>(&sm[(rr8 * 8 + cin) * 72 + (c << 2)]) = v;
        }
    }
    __syncthreads();

    // ---- 8 n-tiles x 4 accumulators ----
    float acc[8][4];
    #pragma unroll
    for (int t = 0; t < 8; ++t) {
        acc[t][0] = acc[t][1] = acc[t][2] = acc[t][3] = 0.f;
    }

    const float4* wp = wfrag4 + warp * 32 + lane;

    #pragma unroll 3
    for (int k = 0; k < 27; ++k) {
        const int dz = k / 9;
        const int dy = (k % 9) / 3;
        const int dx = k % 3;
        const float4 a4 = __ldg(wp + k * 128);
        const unsigned a0 = __float_as_uint(a4.x);
        const unsigned a1 = __float_as_uint(a4.y);
        const unsigned a2 = __float_as_uint(a4.z);
        const unsigned a3 = __float_as_uint(a4.w);
        const int r = dz * 3 + dy;
        const float* rowc = &sm[(r * 8 + (lane & 3)) * 72 + dx + (lane >> 2)];
        const float* rowc4 = rowc + 4 * 72;
        #pragma unroll
        for (int t = 0; t < 8; ++t) {
            const unsigned b0 = __float_as_uint(rowc[t << 3]);
            const unsigned b1 = __float_as_uint(rowc4[t << 3]);
            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[t][0]), "+f"(acc[t][1]),
                  "+f"(acc[t][2]), "+f"(acc[t][3])
                : "r"(a0), "r"(a1), "r"(a2), "r"(a3),
                  "r"(b0), "r"(b1));
        }
    }

    // ---- fused epilogue + NCDHW store ----
    const int basec = warp * 16 + (lane >> 2);
    const int wv = w0 + 2 * (lane & 3);
    const float b0v = __ldg(bias + basec);
    const float s0v = __ldg(sumt + basec);
    const float b8v = __ldg(bias + basec + 8);
    const float s8v = __ldg(sumt + basec + 8);
    float* rowp = out + nOutOff + (long)basec * CPlane + rowOff + wv;
    const long stride8 = 8 * CPlane;

    #pragma unroll
    for (int t = 0; t < 8; ++t) {
        const int w = wv + (t << 3);
        if (w < Wo) {
            const float r0 = epilogue_fn(acc[t][0], b0v, s0v);
            const float r2 = epilogue_fn(acc[t][2], b8v, s8v);
            float* p0 = rowp + (t << 3);
            float* p8 = p0 + stride8;
            if (w + 1 < Wo) {
                const float r1 = epilogue_fn(acc[t][1], b0v, s0v);
                const float r3 = epilogue_fn(acc[t][3], b8v, s8v);
                *reinterpret_cast<float2*>(p0) = make_float2(r0, r1);
                *reinterpret_cast<float2*>(p8) = make_float2(r2, r3);
            } else {
                *p0 = r0;
                *p8 = r2;
            }
        }
    }
}

// ===========================================================================
// v4 fallback kernels (layout conversion + fused epilogue for the cuDNN path)
// ===========================================================================

// -------------- NCDHW -> NDHWC: per-sample [8, M] -> [M, 8] -----------------
__global__ void __launch_bounds__(256)
ncdhw_to_ndhwc_c8_kernel(const float4* __restrict__ in,
                         float4* __restrict__ out,
                         long M) {
    __shared__ float smem[8][260];
    long n = blockIdx.y;
    long m0 = (long)blockIdx.x * 256;
    long M4 = M >> 2;
    const float4* inb = in + n * (8 * M4);
    float4* outb = out + n * (2 * M);
    int t = threadIdx.x;

    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int c = f4 >> 6;
        int m4 = f4 & 63;
        long m = m0 + (m4 << 2);
        if (m < M) {
            float4 v = inb[c * M4 + (m >> 2)];
            *reinterpret_cast<float4*>(&smem[c][m4 << 2]) = v;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int ml = f4 >> 1;
        int half = f4 & 1;
        long m = m0 + ml;
        if (m < M) {
            int c0 = half << 2;
            float4 v;
            v.x = smem[c0 + 0][ml];
            v.y = smem[c0 + 1][ml];
            v.z = smem[c0 + 2][ml];
            v.w = smem[c0 + 3][ml];
            outb[(m << 1) + half] = v;
        }
    }
}

// ------- fused epilogue + transpose: NDHWC [M, C] -> NCDHW [C, M] -----------
__global__ void __launch_bounds__(256)
epilogue_ndhwc_to_ncdhw_kernel(const float4* __restrict__ in,
                               float4* __restrict__ out,
                               const float* __restrict__ bias,
                               const float* __restrict__ sumt,
                               long M, int C) {
    __shared__ float smem[64][33];
    long n = blockIdx.y;
    long m0 = (long)blockIdx.x * 32;
    long M4 = M >> 2;
    int C4 = C >> 2;
    const float4* inb = in + n * (M4 * C);
    float4* outb = out + n * (M4 * C);
    int t = threadIdx.x;

    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int ml = f4 >> 4;
        int c4 = f4 & 15;
        long m = m0 + ml;
        if (c4 < C4 && m < M) {
            float4 v = inb[m * C4 + c4];
            int c0 = c4 << 2;
            v.x = epilogue_fn(v.x, __ldg(bias + c0 + 0), __ldg(sumt + c0 + 0));
            v.y = epilogue_fn(v.y, __ldg(bias + c0 + 1), __ldg(sumt + c0 + 1));
            v.z = epilogue_fn(v.z, __ldg(bias + c0 + 2), __ldg(sumt + c0 + 2));
            v.w = epilogue_fn(v.w, __ldg(bias + c0 + 3), __ldg(sumt + c0 + 3));
            smem[c0 + 0][ml] = v.x;
            smem[c0 + 1][ml] = v.y;
            smem[c0 + 2][ml] = v.z;
            smem[c0 + 3][ml] = v.w;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int c = f4 >> 3;
        int ml = (f4 & 7) << 2;
        long m = m0 + ml;
        if (c < C && m < M) {
            float4 v;
            v.x = smem[c][ml + 0];
            v.y = smem[c][ml + 1];
            v.z = smem[c][ml + 2];
            v.w = smem[c][ml + 3];
            outb[(long)c * M4 + (m >> 2)] = v;
        }
    }
}

__global__ void epilogue_ndhwc_scalar_kernel(const float* __restrict__ in,
                                             const float* __restrict__ bias,
                                             const float* __restrict__ sumt,
                                             float* __restrict__ out,
                                             long n, int C) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        int c = (int)(i % C);
        out[i] = epilogue_fn(in[i], bias[c], sumt[c]);
    }
}

// -------------------------------- host side --------------------------------
torch::Tensor fused_conv3d_epilogue_cuda(torch::Tensor x, torch::Tensor wfrag,
                                         torch::Tensor bias, torch::Tensor sum_t) {
    TORCH_CHECK(x.is_cuda() && wfrag.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    long N = xc.size(0);
    long Di = xc.size(2), Hi = xc.size(3), Wi = xc.size(4);
    long Do = Di - 2, Ho = Hi - 2, Wo = Wi - 2;
    TORCH_CHECK(N >= 1 && Do >= 1 && Ho >= 1 && Wo >= 1, "bad shape");
    TORCH_CHECK(Wo <= 64 && (Wi % 4) == 0, "unsupported width");
    TORCH_CHECK(N * Do <= 65535, "too many z blocks");
    auto out = at::empty({N, 64, Do, Ho, Wo}, xc.options());
    if (out.numel() == 0) return out;
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();
    auto wf = wfrag.contiguous();
    dim3 grid((unsigned)((Wo + 63) / 64), (unsigned)Ho, (unsigned)(N * Do));
    auto stream = at::cuda::getCurrentCUDAStream();
    conv3d_cin8_cout64_tf32_fused_kernel<<<grid, 128, 0, stream>>>(
        xc.data_ptr<float>(),
        reinterpret_cast<const float4*>(wf.data_ptr<float>()),
        bc.data_ptr<float>(), sc.data_ptr<float>(), out.data_ptr<float>(),
        (int)Do, (int)Ho, (int)Wo, (int)Di, (int)Hi, (int)Wi);
    return out;
}

torch::Tensor ncdhw_to_ndhwc_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.dim() == 5, "x must be 5D");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    long N = xc.size(0), C = xc.size(1);
    long M = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(C == 8 && (M & 3) == 0 && N <= 65535, "unsupported shape");
    auto out = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    if (N * M == 0) return out;
    dim3 grid((unsigned)((M + 255) / 256), (unsigned)N);
    auto stream = at::cuda::getCurrentCUDAStream();
    ncdhw_to_ndhwc_c8_kernel<<<grid, 256, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()), M);
    return out;
}

torch::Tensor fused_epilogue_cuda(torch::Tensor y, torch::Tensor bias, torch::Tensor sum_t) {
    TORCH_CHECK(y.is_cuda(), "y must be CUDA");
    TORCH_CHECK(y.dim() == 5, "y must be 5D");
    TORCH_CHECK(y.scalar_type() == at::kFloat, "y must be float32");
    auto yc = y.is_contiguous(at::MemoryFormat::ChannelsLast3d)
                  ? y
                  : y.contiguous(at::MemoryFormat::ChannelsLast3d);
    long N = yc.size(0), C = yc.size(1);
    long M = yc.size(2) * yc.size(3) * yc.size(4);
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();
    auto stream = at::cuda::getCurrentCUDAStream();

    if (N > 0 && M > 0 && (C & 3) == 0 && C <= 64 && (M & 3) == 0 && N <= 65535) {
        auto out = at::empty({yc.size(0), yc.size(1), yc.size(2), yc.size(3), yc.size(4)},
                             yc.options());
        dim3 grid((unsigned)((M + 31) / 32), (unsigned)N);
        epilogue_ndhwc_to_ncdhw_kernel<<<grid, 256, 0, stream>>>(
            reinterpret_cast<const float4*>(yc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            bc.data_ptr<float>(), sc.data_ptr<float>(), M, (int)C);
        return out;
    }
    auto out = at::empty_like(yc);
    long n = out.numel();
    if (n == 0) return out.contiguous();
    long blocks = (n + 255) / 256;
    epilogue_ndhwc_scalar_kernel<<<(int)blocks, 256, 0, stream>>>(
        yc.data_ptr<float>(), bc.data_ptr<float>(), sc.data_ptr<float>(),
        out.data_ptr<float>(), n, (int)C);
    return out.contiguous();
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_conv3d_epilogue_cuda(torch::Tensor x, torch::Tensor wfrag, torch::Tensor bias, torch::Tensor sum_t);\n"
    "torch::Tensor ncdhw_to_ndhwc_cuda(torch::Tensor x);\n"
    "torch::Tensor fused_epilogue_cuda(torch::Tensor y, torch::Tensor bias, torch::Tensor sum_t);"
)

_fused_ext = load_inline(
    name="fused_conv3d_epilogue_v6",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_conv3d_epilogue_cuda", "ncdhw_to_ndhwc_cuda", "fused_epilogue_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)

# Constant weight-fragment packing tables (m16n8k8 A-fragment layout).
# rows[l] = l >> 2 (cout within 16-slice), cols[l] = l & 3 (cin within 4).
_wfrag_rows = (torch.arange(32, dtype=torch.long) >> 2)
_wfrag_cols = (torch.arange(32, dtype=torch.long) & 3)


def _pack_wfrag(w):
    """Pack conv weight [64, 8, 3, 3, 3] into [27, 4, 32, 4] MMA fragments."""
    wt = w.detach().reshape(64, 8, 27).permute(2, 0, 1).contiguous()  # [tap][cout][cin]
    dev = wt.device
    rows = _wfrag_rows.to(dev)  # [32]
    cols = _wfrag_cols.to(dev)  # [32]
    sbase = (torch.arange(4, device=dev, dtype=torch.long) * 16).view(4, 1)
    r0 = (sbase + rows.view(1, 32))   # [4, 32]
    r8 = r0 + 8
    c0 = cols.view(1, 32).expand(4, 32)
    c4 = c0 + 4
    a0 = wt[:, r0, c0]   # [27, 4, 32]
    a1 = wt[:, r8, c0]
    a2 = wt[:, r0, c4]
    a3 = wt[:, r8, c4]
    return torch.stack([a0, a1, a2, a3], dim=-1).contiguous()


class ModelNew(nn.Module):
    """
    Optimized model. Fast path: one custom TF32 tensor-core kernel computes
    the 3x3x3 conv (Cin=8, Cout=64) directly from NCDHW input and writes the
    final NCDHW output with bias + LeakyReLU + tensor-sum + clamp + GELU
    fused in (no intermediates in global memory). Fallback path: cuDNN
    channels_last_3d conv + custom fused epilogue kernels (identical math).
    """
    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape))
        self.ext = _fused_ext
        self._wfrag = None
        self._wfrag_version = -1
        self._ks = kernel_size if isinstance(kernel_size, tuple) else (
            kernel_size, kernel_size, kernel_size
        )
        # Keep the weight in channels_last_3d for the cuDNN fallback path.
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.to(
                memory_format=torch.channels_last_3d
            )

    def _fallback(self, x):
        w = self.conv.weight
        if not w.is_contiguous(memory_format=torch.channels_last_3d):
            with torch.no_grad():
                self.conv.weight.data = w.contiguous(
                    memory_format=torch.channels_last_3d
                )
            w = self.conv.weight
        M_in = x.size(2) * x.size(3) * x.size(4)
        if x.size(1) == 8 and (M_in % 4) == 0 and x.size(0) <= 65535:
            x_cl = self.ext.ncdhw_to_ndhwc_cuda(x)
        else:
            x_cl = x.to(memory_format=torch.channels_last_3d)
        y = F.conv3d(
            x_cl,
            w,
            None,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
        )
        return self.ext.fused_epilogue_cuda(y, self.conv.bias, self.sum_tensor)

    def forward(self, x):
        w = self.conv.weight
        # Fully fused fast path: single custom tensor-core kernel.
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and self._ks == (3, 3, 3)
            and self.conv.in_channels == 8
            and self.conv.out_channels == 64
            and self.conv.stride == (1, 1, 1)
            and self.conv.padding == (0, 0, 0)
            and self.conv.dilation == (1, 1, 1)
            and self.conv.groups == 1
            and self.sum_tensor.numel() == 64
            and x.size(2) >= 3
            and x.size(3) >= 3
            and x.size(4) >= 3
            and (x.size(4) - 2) <= 64
            and (x.size(4) % 4) == 0
            and x.size(0) * (x.size(2) - 2) <= 65535
        ):
            if self._wfrag is None or self._wfrag_version != w._version:
                with torch.no_grad():
                    self._wfrag = _pack_wfrag(w)
                self._wfrag_version = w._version
            return self.ext.fused_conv3d_epilogue_cuda(
                x, self._wfrag, self.conv.bias, self.sum_tensor
            )

        # Reference ops on CPU / odd dtypes.
        if not x.is_cuda or x.dtype != torch.float32 or x.dim() != 5:
            y = self.conv(x)
            y = F.leaky_relu(y, negative_slope=0.2)
            y = y + self.sum_tensor
            y = torch.clamp(y, min=-1.0, max=1.0)
            return F.gelu(y)
        return self._fallback(x)
```