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

# v3 improvements over v2 (which was correct, 2.57x):
#   * The two custom kernels (input layout transform + fused epilogue with
#     output transpose) were occupancy-limited (~37%) because they used a
#     whole-row C*W shared-memory tile. They are replaced by one generic
#     32x32-tile shared-memory transpose kernel (padded to 33 -> bank-conflict
#     free, ~4KB smem/block -> 100% theoretical occupancy), used in both
#     directions:
#       - NCDHW -> NDHWC input conversion for the NHWC tensor-core conv.
#       - NDHWC -> NCDHW fused epilogue applying
#         bias + LeakyReLU(0.2) + per-channel sum + clamp(-1,1) + exact GELU.
#   * Fully coalesced 128B global accesses on both sides of every kernel.
#   * 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 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
}

// Generic tiled transpose over (C x W) row-slabs of a 5D tensor, with an
// optional fused per-channel epilogue applied while loading.
// A row is indexed by (n, rem) where rem enumerates the (d, h) plane.
// Element (c, w) of a row lives at:
//     base + c*cs + w*ws      (base = n*strideN + rem*strideRem)
// NCDHW layout:  cs = DHW, ws = 1
// NDHWC layout:  cs = 1,   ws = C
// grid.x = rows (N*D*H), grid.y = ceil(C/32)*ceil(W/32), block = 256 threads.
__global__ void __launch_bounds__(256)
row_transpose_kernel(const float* __restrict__ in,
                     float* __restrict__ out,
                     const float* __restrict__ bias,   // may be null
                     const float* __restrict__ sumt,   // may be null
                     long strideN_in, long strideRem_in, long cs_in, long ws_in,
                     long strideN_out, long strideRem_out, long cs_out, long ws_out,
                     int C, int W, int DH, int wtiles) {
    __shared__ float smem[32][33];
    long row = blockIdx.x;
    int n = (int)(row / DH);
    long rem = row - (long)n * DH;
    int tileW = (int)(blockIdx.y % wtiles);
    int tileC = (int)(blockIdx.y / wtiles);
    int c0 = tileC << 5;
    int w0 = tileW << 5;
    const float* inb = in + (long)n * strideN_in + rem * strideRem_in;
    float* outb = out + (long)n * strideN_out + rem * strideRem_out;
    int t = threadIdx.x;
    const bool do_epi = (bias != nullptr);

    // ---- load tile into shared memory (apply epilogue on the fly) ---------
    if (ws_in == 1) {            // input contiguous along w (NCDHW)
        int w = t & 31;
        int c = t >> 5;
        int ww = w0 + w;
        bool wok = (ww < W);
        #pragma unroll
        for (int k = 0; k < 4; ++k, c += 8) {
            int cc = c0 + c;
            if (cc < C && wok) {
                float v = inb[(long)cc * cs_in + ww];
                if (do_epi) v = epilogue(v, bias[cc], sumt[cc]);
                smem[c][w] = v;
            }
        }
    } else {                     // input contiguous along c (NDHWC)
        int c = t & 31;
        int w = t >> 5;
        int cc = c0 + c;
        bool cok = (cc < C);
        float b = 0.0f, s = 0.0f;
        if (do_epi && cok) { b = bias[cc]; s = sumt[cc]; }
        #pragma unroll
        for (int k = 0; k < 4; ++k, w += 8) {
            int ww = w0 + w;
            if (cok && ww < W) {
                float v = inb[(long)ww * ws_in + cc];
                if (do_epi) v = epilogue(v, b, s);
                smem[c][w] = v;
            }
        }
    }
    __syncthreads();

    // ---- store tile to global memory in the transposed layout -------------
    if (ws_out == 1) {           // output contiguous along w (NCDHW)
        int w = t & 31;
        int c = t >> 5;
        int ww = w0 + w;
        bool wok = (ww < W);
        #pragma unroll
        for (int k = 0; k < 4; ++k, c += 8) {
            int cc = c0 + c;
            if (cc < C && wok) {
                outb[(long)cc * cs_out + ww] = smem[c][w];
            }
        }
    } else {                     // output contiguous along c (NDHWC)
        int c = t & 31;
        int w = t >> 5;
        int cc = c0 + c;
        bool cok = (cc < C);
        #pragma unroll
        for (int k = 0; k < 4; ++k, w += 8) {
            int ww = w0 + w;
            if (cok && ww < W) {
                outb[(long)ww * ws_out + cc] = smem[c][w];
            }
        }
    }
}

// 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 rows = N * D * H;
    if (rows == 0) return out;
    long DH = D * H, DHW = DH * W;
    int wtiles = (int)((W + 31) / 32), ctiles = (int)((C + 31) / 32);
    dim3 grid((unsigned)rows, (unsigned)(wtiles * ctiles));
    auto stream = at::cuda::getCurrentCUDAStream();
    row_transpose_kernel<<<grid, 256, 0, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(),
        nullptr, nullptr,
        C * DHW, W, DHW, 1,          // input NCDHW
        DHW * C, W * C, 1, C,        // output NDHWC
        (int)C, (int)W, (int)DH, wtiles);
    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 rows = N * D * H;
    long DH = D * H, DHW = DH * W;
    auto stream = at::cuda::getCurrentCUDAStream();

    // Fast path: tiled transpose epilogue (NDHWC -> NCDHW).
    if (rows > 0 && C <= (1 << 22)) {
        auto out = at::empty({N, C, D, H, W}, yc.options());  // NCDHW contiguous
        int wtiles = (int)((W + 31) / 32), ctiles = (int)((C + 31) / 32);
        dim3 grid((unsigned)rows, (unsigned)(wtiles * ctiles));
        row_transpose_kernel<<<grid, 256, 0, stream>>>(
            yc.data_ptr<float>(), out.data_ptr<float>(),
            bc.data_ptr<float>(), sc.data_ptr<float>(),
            DHW * C, W * C, 1, C,    // input NDHWC
            C * DHW, W, DHW, 1,      // output NCDHW
            (int)C, (int)W, (int)DH, wtiles);
        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_v3",
    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 tiled-transpose custom kernel,
      * cuDNN conv3d runs natively in channels_last_3d (no transform kernels
        on the activations),
      * one fused custom kernel applies bias + LeakyReLU + tensor-sum + clamp +
        GELU and writes the final NCDHW-contiguous result via a bank-conflict
        free shared-memory 32x32 tile transpose at full occupancy.
    """
    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.
        if x.is_cuda and x.dtype == torch.float32:
            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)
```