```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 3D convolution algorithm for the fixed shape.
torch.backends.cudnn.benchmark = True

# Custom CUDA extension:
#  - calls the highly-optimized cuDNN 3D convolution
#  - fuses the entire pointwise chain (LeakyReLU -> +sum_tensor -> clamp -> GELU)
#    into a single custom CUDA kernel, eliminating 3 extra full-tensor
#    read/write passes over the large (128, 64, 14, 62, 62) activation.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

// Fused elementwise chain:
//   v = leaky_relu(v, 0.2)
//   v = v + sum_tensor[c]
//   v = clamp(v, -1, 1)
//   v = gelu(v)            (exact erf formulation, matches F.gelu default)
__device__ __forceinline__ float leaky_add_clamp_gelu(float v, float s) {
    v = v > 0.0f ? v : 0.2f * v;
    v = v + s;
    v = fminf(fmaxf(v, -1.0f), 1.0f);
    return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));
}

// Scalar fallback kernel (any spatial size).
__global__ void fused_epilogue_scalar_kernel(float* __restrict__ data,
                                             const float* __restrict__ sum_t,
                                             int spatial, int C) {
    int plane = blockIdx.y;          // plane = n * C + c
    int c = plane % C;
    float s = sum_t[c];
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < spatial) {
        long off = (long)plane * spatial + idx;
        data[off] = leaky_add_clamp_gelu(data[off], s);
    }
}

// Vectorized (float4) kernel used when spatial size is divisible by 4.
__global__ void fused_epilogue_vec4_kernel(float4* __restrict__ data,
                                           const float* __restrict__ sum_t,
                                           int spatial4, int C) {
    int plane = blockIdx.y;          // plane = n * C + c
    int c = plane % C;
    float s = sum_t[c];
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < spatial4) {
        long off = (long)plane * spatial4 + idx;
        float4 v = data[off];
        v.x = leaky_add_clamp_gelu(v.x, s);
        v.y = leaky_add_clamp_gelu(v.y, s);
        v.z = leaky_add_clamp_gelu(v.z, s);
        v.w = leaky_add_clamp_gelu(v.w, s);
        data[off] = v;
    }
}

} // namespace

torch::Tensor conv3d_leaky_add_clamp_gelu(torch::Tensor x,
                                          torch::Tensor weight,
                                          torch::Tensor bias,
                                          torch::Tensor sum_tensor) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && sum_tensor.is_cuda(),
                "All tensors must be CUDA tensors.");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Only float32 is supported.");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto sum_c = sum_tensor.contiguous();

    // cuDNN 3D convolution (stride=1, pad=0, dilation=1, groups=1).
    auto out = at::conv3d(xc, wc, bias);
    TORCH_CHECK(out.is_contiguous(), "Conv output must be contiguous.");

    const int C = (int)out.size(1);
    long spatial_l = 1;
    for (int i = 2; i < out.dim(); ++i) spatial_l *= out.size(i);
    const int nc = (int)(out.size(0) * C);
    TORCH_CHECK(nc <= 65535, "Too many channel planes.");

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

    if ((spatial_l & 3) == 0) {
        int spatial4 = (int)(spatial_l >> 2);
        dim3 grid((spatial4 + threads - 1) / threads, nc);
        fused_epilogue_vec4_kernel<<<grid, threads, 0, stream>>>(
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            sum_c.data_ptr<float>(), spatial4, C);
    } else {
        int spatial = (int)spatial_l;
        dim3 grid((spatial + threads - 1) / threads, nc);
        fused_epilogue_scalar_kernel<<<grid, threads, 0, stream>>>(
            out.data_ptr<float>(), sum_c.data_ptr<float>(), spatial, C);
    }

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor conv3d_leaky_add_clamp_gelu(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, torch::Tensor sum_tensor);"
)

fused_module = load_inline(
    name="conv3d_fused_epilogue",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["conv3d_leaky_add_clamp_gelu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN Conv3d followed by a single fused custom CUDA kernel
    that applies LeakyReLU, the per-channel sum, clamp, and GELU in one pass.
    """

    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.fused_op = fused_module

    def forward(self, x):
        return self.fused_op.conv3d_leaky_add_clamp_gelu(
            x, self.conv.weight, self.conv.bias, self.sum_tensor
        )
```