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

# Custom CUDA kernel that fuses the four pointwise operators that follow the
# convolution (LeakyReLU -> add per-channel sum tensor -> clamp -> exact GELU)
# into a single memory-bound pass over the convolution output.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr float kNegSlope = 0.2f;
constexpr float kInvSqrt2 = 0.70710678118654752440f;

__device__ __forceinline__ float epilogue(float v, float s) {
    // LeakyReLU(negative_slope = 0.2)
    v = v > 0.0f ? v : kNegSlope * v;
    // + per-channel sum tensor
    v = v + s;
    // clamp(-1, 1)
    v = fminf(fmaxf(v, -1.0f), 1.0f);
    // exact (erf) GELU
    return 0.5f * v * (1.0f + erff(v * kInvSqrt2));
}

// Vectorized path: each thread handles one float4 (16 bytes in / 16 bytes out).
// Requires numel % 4 == 0 and per-channel spatial size % 4 == 0 so that a
// float4 never straddles a channel boundary.
__global__ void fused_epilogue_vec4_kernel(
        const float4* __restrict__ in,
        const float* __restrict__ sum,
        float4* __restrict__ out,
        const long nvec,
        const long dhw_vec,
        const int channels) {
    const long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i >= nvec) return;
    const int c = (int)((i / dhw_vec) % channels);
    const float s = sum[c];
    float4 v = in[i];
    v.x = epilogue(v.x, s);
    v.y = epilogue(v.y, s);
    v.z = epilogue(v.z, s);
    v.w = epilogue(v.w, s);
    out[i] = v;
}

// Generic scalar fallback path for arbitrary shapes.
__global__ void fused_epilogue_scalar_kernel(
        const float* __restrict__ in,
        const float* __restrict__ sum,
        float* __restrict__ out,
        const long n,
        const long dhw,
        const int channels) {
    const long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i >= n) return;
    const int c = (int)((i / dhw) % channels);
    out[i] = epilogue(in[i], sum[c]);
}

} // namespace

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

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

    const int channels = (xc.dim() >= 2) ? (int)xc.size(1) : 1;
    long dhw = 1;
    for (int d = 2; d < xc.dim(); ++d) dhw *= xc.size(d);
    TORCH_CHECK(sc.numel() == channels, "sum tensor must have one value per channel");

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

    if ((numel & 3) == 0 && (dhw & 3) == 0) {
        const long nvec = numel >> 2;
        const long dhw_vec = dhw >> 2;
        const unsigned int blocks = (unsigned int)((nvec + threads - 1) / threads);
        fused_epilogue_vec4_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            sc.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            nvec, dhw_vec, channels);
    } else {
        const unsigned int blocks = (unsigned int)((numel + threads - 1) / threads);
        fused_epilogue_scalar_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), sc.data_ptr<float>(), out.data_ptr<float>(),
            numel, dhw, channels);
    }
    return out;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized version of Model: keeps the cuDNN Conv3d (autotuned) and replaces
    the four trailing pointwise operators with a single fused CUDA kernel,
    eliminating multiple full-tensor read/write passes.
    """
    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))
        # Let cuDNN autotune the fastest conv3d algorithm/engine and allow TF32.
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True

    def forward(self, x):
        x = F.conv3d(x, self.conv.weight, self.conv.bias)
        x = fused_epilogue_ext.fused_epilogue(x, self.sum_tensor)
        return x
```