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

# Fused epilogue kernel: (conv bias) + LeakyReLU(0.2) + add per-channel
# sum_tensor + clamp(-1, 1) + exact GELU, all in a single memory pass.
fused_epilogue_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 sum tensor
    v = fminf(fmaxf(v, CLAMP_LO), CLAMP_HI);       // clamp(-1, 1)
    return 0.5f * v * (1.0f + erff(v * RSQRT2));   // exact GELU
}

__global__ void fused_epilogue_vec4_kernel(
    const float4* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ sum_t,
    float4* __restrict__ out,
    long n_vec,
    int channels,
    long spatial_vec) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n_vec) {
        int c = (int)((i / spatial_vec) % channels);
        float b = bias[c];
        float s = sum_t[c];
        float4 v = x[i];
        v.x = epilogue(v.x, b, s);
        v.y = epilogue(v.y, b, s);
        v.z = epilogue(v.z, b, s);
        v.w = epilogue(v.w, b, s);
        out[i] = v;
    }
}

__global__ void fused_epilogue_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ sum_t,
    float* __restrict__ out,
    long n,
    int channels,
    long spatial) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        int c = (int)((i / spatial) % channels);
        out[i] = epilogue(x[i], bias[c], sum_t[c]);
    }
}

torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor 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 xc = x.contiguous();
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();
    auto out = torch::empty_like(xc);

    long n = xc.numel();
    if (n == 0) return out;

    int C = (int)xc.size(1);
    long spatial = 1;
    for (int i = 2; i < (int)xc.dim(); ++i) spatial *= xc.size(i);

    const int threads = 256;
    const float* xp = xc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    const float* sp = sc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    if ((n & 3L) == 0 && (spatial & 3L) == 0) {
        // Vectorized path: 4 floats per thread (16 bytes per load/store).
        long n_vec = n >> 2;
        long spatial_vec = spatial >> 2;
        long blocks = (n_vec + threads - 1) / threads;
        fused_epilogue_vec4_kernel<<<blocks, threads>>>(
            reinterpret_cast<const float4*>(xp), bp, sp,
            reinterpret_cast<float4*>(op), n_vec, C, spatial_vec);
    } else {
        long blocks = (n + threads - 1) / threads;
        fused_epilogue_scalar_kernel<<<blocks, threads>>>(xp, bp, sp, op, n, C, spatial);
    }
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor bias, torch::Tensor sum_t);"
)

fused_epilogue = load_inline(
    name="fused_conv3d_epilogue",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_epilogue_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN 3D convolution (bias handled in the epilogue) plus a
    single fused CUDA kernel for bias + LeakyReLU + tensor-sum + clamp + GELU.
    This replaces 4 separate elementwise CUDA kernels (and their global-memory
    round trips) with one vectorized 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.epilogue = fused_epilogue

    def forward(self, x):
        # Bias is applied inside the fused epilogue kernel.
        x = F.conv3d(
            x,
            self.conv.weight,
            None,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
        )
        return self.epilogue.fused_epilogue_cuda(x, self.conv.bias, self.sum_tensor)
```