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

# Autotune cuDNN for the fixed conv shapes used by this model.
torch.backends.cudnn.benchmark = True

# Fused CUDA kernel that applies, in a single memory pass:
#   LeakyReLU(0.2) -> + per-channel sum_tensor -> clamp(-1, 1) -> exact GELU
fused_postprocess_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float fused_op(float v, float s) {
    // LeakyReLU with negative_slope = 0.2
    v = (v > 0.0f) ? v : v * 0.2f;
    // Add per-channel bias (sum_tensor)
    v += s;
    // 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. Grid: (spatial blocks, N*C slices).
// Requires spatial_size % 4 == 0 so a float4 never crosses a channel boundary.
__global__ void fused_vec4_kernel(const float4* __restrict__ x,
                                  const float* __restrict__ sum_tensor,
                                  float4* __restrict__ out,
                                  int spatial_vec,
                                  int channels) {
    const int nc = blockIdx.y;
    const float s = sum_tensor[nc % channels];
    const long base = (long)nc * spatial_vec;
    const float4* xs = x + base;
    float4* os = out + base;

    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < spatial_vec; i += stride) {
        float4 v = xs[i];
        float4 r;
        r.x = fused_op(v.x, s);
        r.y = fused_op(v.y, s);
        r.z = fused_op(v.z, s);
        r.w = fused_op(v.w, s);
        os[i] = r;
    }
}

// Scalar fallback path for arbitrary shapes.
__global__ void fused_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ sum_tensor,
                                    float* __restrict__ out,
                                    long spatial_size,
                                    int channels) {
    const int nc = blockIdx.y;
    const float s = sum_tensor[nc % channels];
    const long base = (long)nc * spatial_size;
    const float* xs = x + base;
    float* os = out + base;

    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < spatial_size; i += stride) {
        os[i] = fused_op(xs[i], s);
    }
}

torch::Tensor fused_postprocess_cuda(torch::Tensor x, torch::Tensor sum_tensor) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(sum_tensor.is_cuda(), "sum_tensor must be a CUDA tensor");

    auto xc = x.contiguous();
    auto sc = sum_tensor.contiguous();

    const int channels = xc.size(1);
    long spatial = 1;
    for (int i = 2; i < xc.dim(); ++i) spatial *= xc.size(i);
    const long nc_total = xc.numel() / spatial;

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

    if (spatial % 4 == 0) {
        const int spatial_vec = (int)(spatial >> 2);
        long want = (spatial_vec + block - 1) / block;
        int gx = (int)(want > 65535 ? 65535 : want);
        dim3 grid(gx, (unsigned)nc_total);
        // In-place: each element is read once and written once at the same index.
        fused_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            sc.data_ptr<float>(),
            reinterpret_cast<float4*>(xc.data_ptr<float>()),
            spatial_vec, channels);
    } else {
        long want = (spatial + block - 1) / block;
        int gx = (int)(want > 65535 ? 65535 : want);
        dim3 grid(gx, (unsigned)nc_total);
        fused_scalar_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(),
            sc.data_ptr<float>(),
            xc.data_ptr<float>(),
            spatial, channels);
    }
    return xc;
}
"""

fused_postprocess_cpp_source = (
    "torch::Tensor fused_postprocess_cuda(torch::Tensor x, torch::Tensor sum_tensor);"
)

fused_postprocess = load_inline(
    name="fused_postprocess",
    cpp_sources=fused_postprocess_cpp_source,
    cuda_sources=fused_postprocess_source,
    functions=["fused_postprocess_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN 3D convolution and fuses
    LeakyReLU + tensor addition + clamp + GELU into one 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)
        x = fused_postprocess.fused_postprocess_cuda(
            x, self.sum_tensor.view(-1))
        return x
```