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

# v4: The layout conversions are re-formulated as plain per-sample 2D matrix
# transposes (valid because both NCDHW and NDHWC are contiguous per sample):
#   * input  NCDHW [8, M] -> NDHWC [M, 8]   (M = D*H*W)
#   * output NDHWC [M, 64] -> NCDHW [64, M] with the fused epilogue
#     (bias + LeakyReLU(0.2) + per-channel sum + clamp(-1,1) + exact GELU)
#     applied on the fly while loading.
# Both kernels use float4 global loads AND stores, padded shared-memory tiles
# arranged so all shared accesses are (nearly) bank-conflict free, run at full
# occupancy, and move exactly the minimum required bytes (memory bound).
# cuDNN still runs the conv natively in channels_last_3d.
_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
}

// -------------- NCDHW -> NDHWC: per-sample [8, M] -> [M, 8] -----------------
// Tile: 256 m x 8 c per block, 256 threads, float4 on both global sides.
// Requires C == 8 and M % 4 == 0. grid = (ceil(M/256), N).
__global__ void __launch_bounds__(256)
ncdhw_to_ndhwc_c8_kernel(const float4* __restrict__ in,
                         float4* __restrict__ out,
                         long M) {
    __shared__ float smem[8][260];  // 260 % 4 == 0 -> aligned float4 writes
    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);   // 0..511
        int c = f4 >> 6;          // 0..7 (64 float4 per channel row)
        int m4 = f4 & 63;         // 0..63
        long m = m0 + (m4 << 2);
        if (m < M) {
            float4 v = inb[c * M4 + (m >> 2)];          // 512B coalesced
            *reinterpret_cast<float4*>(&smem[c][m4 << 2]) = v;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);   // 0..511
        int ml = f4 >> 1;         // 0..255
        int half = f4 & 1;
        long m = m0 + ml;
        if (m < M) {
            int c0 = half << 2;                         // 0 or 4
            float4 v;
            v.x = smem[c0 + 0][ml];                     // conflict-free reads
            v.y = smem[c0 + 1][ml];
            v.z = smem[c0 + 2][ml];
            v.w = smem[c0 + 3][ml];
            outb[(m << 1) + half] = v;                  // 512B coalesced
        }
    }
}

// ------- fused epilogue + transpose: NDHWC [M, C] -> NCDHW [C, M] -----------
// Tile: 32 m x 64 c per block, 256 threads, float4 global loads/stores.
// Requires C % 4 == 0, C <= 64, M % 4 == 0. grid = (ceil(M/32), N).
__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];  // +1 pad: conflict-free scalar reads
    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);   // 0..511
        int ml = f4 >> 4;         // 0..31
        int c4 = f4 & 15;         // 0..15
        long m = m0 + ml;
        if (c4 < C4 && m < M) {
            float4 v = inb[m * C4 + c4];                // 256B+ coalesced
            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);       // 0..511
        int c = f4 >> 3;              // 0..63
        int ml = (f4 & 7) << 2;       // 0,4,...,28
        long m = m0 + ml;
        if (c < C && m < M) {
            float4 v;
            v.x = smem[c][ml + 0];    // conflict-free scalar reads
            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;          // coalesced float4
        }
    }
}

// Fallback: scalar elementwise epilogue directly on NDHWC data (any shape).
__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]);
    }
}

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());  // NCDHW contiguous
        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;
    }
    // Generic fallback: elementwise on NDHWC, then convert layout with ATen.
    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 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_nhwc_epilogue_v4",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["ncdhw_to_ndhwc_cuda", "fused_epilogue_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * input converted NCDHW->NDHWC by a fast fully-vectorized custom
        transpose kernel,
      * cuDNN conv3d runs natively in channels_last_3d (tensor cores),
      * one fused custom kernel applies bias + LeakyReLU + tensor-sum + clamp +
        GELU and writes the final NCDHW-contiguous result via a shared-memory
        tile transpose with float4 accesses on both sides.
    """
    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
        # Keep the weight in channels_last_3d so cuDNN avoids weight transforms.
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.to(
                memory_format=torch.channels_last_3d
            )

    def forward(self, x):
        # Safety fallback for unexpected parameter/input configurations.
        if (
            self.sum_tensor.numel() != self.conv.out_channels
            or x.dim() != 5
            or not x.is_cuda
            or x.dtype != torch.float32
        ):
            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)

        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

        # Fast custom NCDHW -> NDHWC conversion (vectorized for C == 8).
        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)

        # Bias is applied inside the fused epilogue kernel.
        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)
```