```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 given shapes.
torch.backends.cudnn.benchmark = True

# Custom CUDA kernel that fuses the entire pointwise epilogue after the Conv3d:
#   LeakyReLU(0.2) -> + per-channel sum_tensor -> clamp(-1, 1) -> exact GELU
# This replaces 4 separate memory-bound elementwise kernels with a single pass.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define GELU_INV_SQRT2 0.7071067811865476f  // 1/sqrt(2)
#define LEAKY_SLOPE    0.2f

__device__ __forceinline__ float gelu_exact(float v) {
    return 0.5f * v * (1.0f + erff(v * GELU_INV_SQRT2));
}

__device__ __forceinline__ float epilogue(float v, float s) {
    v = v > 0.0f ? v : LEAKY_SLOPE * v;    // LeakyReLU(0.2)
    v = v + s;                             // add per-channel sum_tensor
    v = fminf(fmaxf(v, -1.0f), 1.0f);      // clamp to [-1, 1]
    return gelu_exact(v);                  // exact (erf) GELU
}

// Vectorized (float4) path: grid.y indexes the N*C image dimension so the
// channel index is obtained without any integer division. Requires the
// per-image spatial size (D*H*W) to be divisible by 4 so every float4
// lies entirely within one channel plane.
__global__ void fused_epilogue_vec4_kernel(const float4* __restrict__ x,
                                           const float* __restrict__ sum_t,
                                           float4* __restrict__ out,
                                           int sp4, int C) {
    int sp_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (sp_idx >= sp4) return;
    int c = blockIdx.y % C;
    long idx = (long)blockIdx.y * sp4 + sp_idx;
    float s = sum_t[c];
    float4 v = x[idx];
    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[idx] = v;
}

// Generic scalar fallback for arbitrary shapes.
__global__ void fused_epilogue_scalar_kernel(const float* __restrict__ x,
                                             const float* __restrict__ sum_t,
                                             float* __restrict__ out,
                                             long total, long sp, int C) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (int)((idx / sp) % C);
    out[idx] = epilogue(x[idx], sum_t[c]);
}

torch::Tensor fused_epilogue_cuda(torch::Tensor x, torch::Tensor sum_t) {
    TORCH_CHECK(x.is_cuda() && sum_t.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto x_c = x.contiguous();
    auto s_c = sum_t.contiguous();
    auto out = torch::empty_like(x_c);

    long total = x_c.numel();
    if (total == 0) return out;

    int C = (int)x_c.size(1);
    long N = (long)x_c.size(0);
    long sp = 1;
    for (int i = 2; i < x_c.dim(); i++) sp *= x_c.size(i);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (sp % 4 == 0 && N * (long)C <= 65535) {
        int sp4 = (int)(sp / 4);
        dim3 block(256);
        dim3 grid((unsigned)((sp4 + block.x - 1) / block.x), (unsigned)(N * C));
        fused_epilogue_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            s_c.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            sp4, C);
    } else {
        int block = 256;
        long grid = (total + block - 1) / block;
        fused_epilogue_scalar_kernel<<<(unsigned)grid, block, 0, stream>>>(
            x_c.data_ptr<float>(),
            s_c.data_ptr<float>(),
            out.data_ptr<float>(),
            total, sp, C);
    }
    return out;
}
"""

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

fused_epilogue = load_inline(
    name="fused_epilogue_ext",
    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: Conv3d is delegated to cuDNN (with autotuning enabled),
    while LeakyReLU + per-channel tensor addition + clamp + GELU are fused
    into a single custom CUDA kernel, eliminating three full read/write
    passes over the activation tensor.
    """
    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)
```