```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
from torch.utils.cpp_extension import load_inline

# Let cuDNN autotune the best Conv3d algorithm for the fixed input shape.
torch.backends.cudnn.benchmark = True

# Custom fused CUDA kernel for the element-wise epilogue:
#   LeakyReLU(0.2) -> + per-channel bias -> clamp(-1, 1) -> exact GELU
# This replaces 4 separate memory-bound kernels (each reading+writing the full
# activation tensor) with a single vectorized pass.
fused_epilogue_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float fused_op(float v, float b) {
    // LeakyReLU with negative_slope = 0.2
    v = v > 0.0f ? v : 0.2f * v;
    // Broadcast add of the per-channel sum tensor
    v = v + b;
    // Clamp to [-1, 1]
    v = fminf(fmaxf(v, -1.0f), 1.0f);
    // Exact GELU: 0.5 * x * (1 + erf(x / sqrt(2)))
    return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));
}

// Vectorized (float4) path. blockIdx.y indexes the (N*C) slice so the
// channel index is obtained with a cheap modulo instead of a full division.
__global__ void fused_epilogue_vec4_kernel(const float4* __restrict__ x,
                                           const float* __restrict__ bias,
                                           float4* __restrict__ out,
                                           int inner4,
                                           int C) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < inner4) {
        float b = bias[blockIdx.y % C];
        long base = (long)blockIdx.y * (long)inner4 + (long)i;
        float4 v = x[base];
        v.x = fused_op(v.x, b);
        v.y = fused_op(v.y, b);
        v.z = fused_op(v.z, b);
        v.w = fused_op(v.w, b);
        out[base] = v;
    }
}

// Scalar fallback for inner sizes not divisible by 4.
__global__ void fused_epilogue_scalar_kernel(const float* __restrict__ x,
                                             const float* __restrict__ bias,
                                             float* __restrict__ out,
                                             long inner,
                                             int C) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < inner) {
        float b = bias[blockIdx.y % C];
        long base = (long)blockIdx.y * inner + i;
        out[base] = fused_op(x[base], b);
    }
}

torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto x_c = x.contiguous();
    auto bias_c = bias.contiguous();
    auto out = torch::empty_like(x_c);

    const int C = (int)x_c.size(1);
    long inner = 1;
    for (int d = 2; d < x_c.dim(); ++d) inner *= x_c.size(d);
    const long NC = (long)x_c.size(0) * (long)C;

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

    if (inner % 4 == 0) {
        int inner4 = (int)(inner / 4);
        dim3 grid((unsigned)((inner4 + block_size - 1) / block_size), (unsigned)NC);
        fused_epilogue_vec4_kernel<<<grid, block_size, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            bias_c.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            inner4, C);
    } else {
        dim3 grid((unsigned)((inner + block_size - 1) / block_size), (unsigned)NC);
        fused_epilogue_scalar_kernel<<<grid, block_size, 0, stream>>>(
            x_c.data_ptr<float>(),
            bias_c.data_ptr<float>(),
            out.data_ptr<float>(),
            inner, C);
    }
    return out;
}
"""

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

# Compile the inline CUDA epilogue
fused_epilogue = load_inline(
    name="fused_leaky_add_clamp_gelu_epilogue",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_source,
    functions=["fused_epilogue_cuda"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Model that performs a 3D convolution, applies LeakyReLU, sums with a tensor, clamps, and applies GELU activation.

    Optimized version: the Conv3d is executed by cuDNN (with autotuning), while
    LeakyReLU + broadcast-add + clamp + GELU are fused into a single custom
    vectorized CUDA kernel, eliminating 3 full-tensor round trips to memory.
    """
    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))

    def forward(self, x):
        x = self.conv(x)
        return fused_epilogue.fused_epilogue_cuda(x, self.sum_tensor.view(-1))
```