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

# Key idea (v2): keep the convolution fully in channels_last_3d (NDHWC) so
# cuDNN runs its fast NHWC tensor-core kernel with NO layout-transform kernels.
# A tiny custom kernel converts the input NCDHW->NDHWC (replaces cuDNN's slow
# precompiled nchwToNhwc), and a fused epilogue kernel applies
# bias + LeakyReLU(0.2) + per-channel add + clamp(-1,1) + exact GELU while
# transposing NDHWC->NCDHW through shared memory (replaces cuDNN's slow
# nhwcToNchw + 4 separate elementwise kernels) in one memory-bound pass.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define NEG_SLOPE 0.2f
#define CLAMP_LO  -1.0f
#define CLAMP_HI   1.0f
#define RSQRT2    0.70710678118654752440f  // 1 / sqrt(2)

__device__ __forceinline__ float epilogue(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, CLAMP_LO), CLAMP_HI);       // clamp(-1, 1)
    return 0.5f * v * (1.0f + erff(v * RSQRT2));   // exact GELU
}

// ---------------- NCDHW -> NDHWC transform (one block per (n,d,h) tile) ------
__global__ void ncdhw_to_ndhwc_kernel(const float* __restrict__ in,
                                      float* __restrict__ out,
                                      int C, int D, int H, int W) {
    extern __shared__ float smem[];  // (C+1)*W floats, [w*(C+1) + c]
    int DH = D * H;
    long ndh = blockIdx.x;
    long n = ndh / DH;
    int rem = (int)(ndh - n * DH);
    int d = rem / H;
    int h = rem - d * H;
    long HW = (long)H * W;
    long cstride = (long)D * HW;
    const float* in_base = in + (n * C * D + d) * HW + h * W;
    float* out_base = out + ((n * D + d) * H + h) * ((long)W * C);
    int CW = C * W;
    int Cp = C + 1;
    for (int idx = threadIdx.x; idx < CW; idx += blockDim.x) {
        int c = idx / W;
        int w = idx - c * W;
        smem[w * Cp + c] = in_base[c * cstride + w];   // coalesced loads along w
    }
    __syncthreads();
    for (int idx = threadIdx.x; idx < CW; idx += blockDim.x) {
        int w = idx / C;
        int c = idx - w * C;
        out_base[idx] = smem[w * Cp + c];              // fully coalesced stores
    }
}

// ---- fused epilogue: NDHWC in -> NCDHW out (one block per (n,d,h) tile) -----
__global__ void epilogue_ndhwc_to_ncdhw_kernel(const float* __restrict__ in,
                                               const float* __restrict__ bias,
                                               const float* __restrict__ sumt,
                                               float* __restrict__ out,
                                               int C, int D, int H, int W) {
    extern __shared__ float smem[];  // C*(W+1) floats, [c*(W+1) + w]
    int DH = D * H;
    long ndh = blockIdx.x;
    long n = ndh / DH;
    int rem = (int)(ndh - n * DH);
    int d = rem / H;
    int h = rem - d * H;
    long HW = (long)H * W;
    long cstride = (long)D * HW;
    const float* in_base = in + ((n * D + d) * H + h) * ((long)W * C);
    float* out_base = out + (n * C * D + d) * HW + h * W;
    int CW = C * W;
    int Wp = W + 1;
    for (int idx = threadIdx.x; idx < CW; idx += blockDim.x) {
        int w = idx / C;
        int c = idx - w * C;
        float v = epilogue(in_base[idx], bias[c], sumt[c]);  // coalesced loads
        smem[c * Wp + w] = v;
    }
    __syncthreads();
    for (int idx = threadIdx.x; idx < CW; idx += blockDim.x) {
        int c = idx / W;
        int w = idx - c * W;
        out_base[c * cstride + w] = smem[c * Wp + w];        // coalesced stores
    }
}

// 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(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), D = xc.size(2), H = xc.size(3), W = xc.size(4);
    auto out = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    long blocks = N * D * H;
    if (blocks == 0) return out;
    int smem = (int)((C + 1) * W * 4);
    TORCH_CHECK(smem <= 48 * 1024, "C*W too large for transform kernel");
    ncdhw_to_ndhwc_kernel<<<(int)blocks, 256, smem>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(),
        (int)C, (int)D, (int)H, (int)W);
    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), D = yc.size(2), H = yc.size(3), W = yc.size(4);
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();

    long smem_floats = C * (W + 1);
    if (smem_floats * 4 <= 48 * 1024) {
        auto out = at::empty({N, C, D, H, W}, yc.options());  // NCDHW contiguous
        long blocks = N * D * H;
        if (blocks == 0) return out;
        epilogue_ndhwc_to_ncdhw_kernel<<<(int)blocks, 256, (int)(smem_floats * 4)>>>(
            yc.data_ptr<float>(), bc.data_ptr<float>(), sc.data_ptr<float>(),
            out.data_ptr<float>(), (int)C, (int)D, (int)H, (int)W);
        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>>>(
        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_v2",
    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 custom kernel,
      * cuDNN conv3d runs natively in channels_last_3d (no transform kernels),
      * one fused custom kernel applies bias + LeakyReLU + tensor-sum + clamp +
        GELU and writes the final NCDHW-contiguous result via an in-kernel
        shared-memory transpose.
    """
    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 shapes: plain eager path.
        if self.sum_tensor.numel() != self.conv.out_channels 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)

        w = self.conv.weight
        if not w.is_contiguous(memory_format=torch.channels_last_3d):
            w = w.contiguous(memory_format=torch.channels_last_3d)

        # Fast custom NCDHW -> NDHWC conversion (fits in shared memory here).
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and (x.size(1) + 1) * x.size(4) * 4 <= 48 * 1024
        ):
            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)
```