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

# Fused CUDA kernel implementing the whole pointwise chain that follows the
# 3D convolution in a single memory pass:
#   (+ conv bias) -> LeakyReLU(0.2) -> (+ sum_tensor) -> clamp(-1, 1) -> GELU
# The convolution bias is folded into this kernel (conv is called with
# bias=None), which removes one extra full read/write pass over the tensor.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ float epilogue(float v, float cb, float st) {
    v += cb;                                                        // conv bias
    v = v > 0.0f ? v : 0.2f * v;                                    // leaky relu
    v += st;                                                        // sum tensor
    v = fminf(fmaxf(v, -1.0f), 1.0f);                               // clamp
    return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));   // gelu (erf)
}

// NDHWC layout (channels_last_3d): channel is the fastest varying dim.
// One float4 = 4 channels of the same spatial location.
__global__ void fused_epilogue_ndhwc_vec4(const float4* __restrict__ x,
                                          const float4* __restrict__ cb4,
                                          const float4* __restrict__ st4,
                                          float4* __restrict__ out,
                                          int n4, int C4, int c4mask) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n4) return;
    int c4 = (c4mask >= 0) ? (i & c4mask) : (i % C4);
    float4 v  = x[i];
    float4 cb = cb4[c4];
    float4 st = st4[c4];
    v.x = epilogue(v.x, cb.x, st.x);
    v.y = epilogue(v.y, cb.y, st.y);
    v.z = epilogue(v.z, cb.z, st.z);
    v.w = epilogue(v.w, cb.w, st.w);
    out[i] = v;
}

// NCDHW layout: 3D grid = (spatial/4, C, N), no integer division needed.
__global__ void fused_epilogue_ncdhw_vec4(const float4* __restrict__ x,
                                          const float* __restrict__ cb,
                                          const float* __restrict__ st,
                                          float4* __restrict__ out,
                                          int S4) {
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= S4) return;
    int c = blockIdx.y;
    long long slab = ((long long)blockIdx.z * gridDim.y + c) * (long long)S4;
    float cbv = cb[c];
    float stv = st[c];
    float4 v = x[slab + t];
    v.x = epilogue(v.x, cbv, stv);
    v.y = epilogue(v.y, cbv, stv);
    v.z = epilogue(v.z, cbv, stv);
    v.w = epilogue(v.w, cbv, stv);
    out[slab + t] = v;
}

// Scalar NCDHW fallback for arbitrary shapes.
__global__ void fused_epilogue_ncdhw_scalar(const float* __restrict__ x,
                                            const float* __restrict__ cb,
                                            const float* __restrict__ st,
                                            float* __restrict__ out,
                                            int S) {
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= S) return;
    int c = blockIdx.y;
    long long slab = ((long long)blockIdx.z * gridDim.y + c) * (long long)S;
    out[slab + t] = epilogue(x[slab + t], cb[c], st[c]);
}

torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor conv_bias, torch::Tensor sum_t) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto st = sum_t.contiguous();
    torch::Tensor cb = conv_bias.defined() ? conv_bias.contiguous()
                                           : torch::zeros_like(st);

    if (!x.is_contiguous() && !x.is_contiguous(at::MemoryFormat::ChannelsLast3d)) {
        x = x.contiguous();
    }
    auto out = torch::empty_like(x);

    const int64_t N = x.size(0);
    const int64_t C = x.size(1);
    int64_t S = 1;
    for (int i = 2; i < x.dim(); ++i) S *= x.size(i);

    const float* xp  = x.data_ptr<float>();
    const float* cbp = cb.data_ptr<float>();
    const float* stp = st.data_ptr<float>();
    float* op = out.data_ptr<float>();

    auto stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    if (x.is_contiguous(at::MemoryFormat::ChannelsLast3d) &&
        (C % 4 == 0) && (x.numel() / 4) <= (int64_t)INT32_MAX) {
        int64_t n4 = x.numel() / 4;
        int C4 = (int)(C / 4);
        int mask = ((C4 & (C4 - 1)) == 0) ? (C4 - 1) : -1;
        int64_t blocks = (n4 + threads - 1) / threads;
        fused_epilogue_ndhwc_vec4<<<(unsigned int)blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xp),
            reinterpret_cast<const float4*>(cbp),
            reinterpret_cast<const float4*>(stp),
            reinterpret_cast<float4*>(op),
            (int)n4, C4, mask);
    } else {
        TORCH_CHECK(N <= 65535 && C <= 65535, "N/C too large for 3D grid");
        if ((S % 4) == 0 && (S / 4) <= (int64_t)INT32_MAX) {
            int S4 = (int)(S / 4);
            dim3 grid((unsigned int)((S4 + threads - 1) / threads),
                      (unsigned int)C, (unsigned int)N);
            fused_epilogue_ncdhw_vec4<<<grid, threads, 0, stream>>>(
                reinterpret_cast<const float4*>(xp), cbp, stp,
                reinterpret_cast<float4*>(op), S4);
        } else {
            TORCH_CHECK(S <= (int64_t)INT32_MAX, "spatial size too large");
            dim3 grid((unsigned int)((S + threads - 1) / threads),
                      (unsigned int)C, (unsigned int)N);
            fused_epilogue_ncdhw_scalar<<<grid, threads, 0, stream>>>(
                xp, cbp, stp, op, (int)S);
        }
    }
    return out;
}
"""

cpp_source = "torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor conv_bias, torch::Tensor sum_t);"

_fused_epilogue = load_inline(
    name="fused_conv3d_epilogue",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_epilogue_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version:
      * conv3d runs in channels_last_3d (NDHWC) so cuDNN can use its fastest
        (tensor-core) kernels; cuDNN benchmark mode autotunes the algorithm.
      * conv bias + LeakyReLU + sum_tensor + clamp + GELU are all fused into a
        single custom CUDA kernel (one memory pass instead of five).
    """
    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape))
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.contiguous(
                memory_format=torch.channels_last_3d)

    def forward(self, x):
        if x.dim() == 5 and not x.is_contiguous(memory_format=torch.channels_last_3d):
            x = x.contiguous(memory_format=torch.channels_last_3d)
        if not self.conv.weight.is_contiguous(memory_format=torch.channels_last_3d):
            with torch.no_grad():
                self.conv.weight.data = self.conv.weight.data.contiguous(
                    memory_format=torch.channels_last_3d)
        # bias is applied inside the fused epilogue kernel
        y = F.conv3d(x, self.conv.weight, None, self.conv.stride,
                     self.conv.padding, self.conv.dilation, self.conv.groups)
        bias = self.conv.bias
        if bias is None:
            bias = torch.zeros(self.conv.out_channels, device=y.device, dtype=y.dtype)
        return _fused_epilogue.fused_epilogue_cuda(y, bias, self.sum_tensor)
```