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

# Fused CUDA kernel implementing the whole post-convolution chain:
#   LeakyReLU(0.2) -> + per-channel sum_tensor -> clamp(-1, 1) -> exact GELU
# This collapses 4 memory-bound elementwise kernels (and their intermediates)
# into a single read/write pass over the convolution output.
fused_post_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <climits>

namespace {

__device__ __forceinline__ float fused_chain(float v, float s) {
    // LeakyReLU with negative_slope = 0.2
    v = (v > 0.0f) ? v : (v * 0.2f);
    // broadcast per-channel sum
    v = v + s;
    // clamp to [-1, 1]
    v = fminf(fmaxf(v, -1.0f), 1.0f);
    // exact (erf-based) GELU: 0.5 * x * (1 + erf(x / sqrt(2)))
    return 0.5f * v * (1.0f + erff(v / 1.4142135623730951f));
}

template <typename IndexT>
__global__ void fused_chain_scalar_kernel(const float* __restrict__ in,
                                          float* __restrict__ out,
                                          const float* __restrict__ sum_t,
                                          int C, IndexT spatial, IndexT total) {
    IndexT idx = (IndexT)blockIdx.x * (IndexT)blockDim.x + (IndexT)threadIdx.x;
    if (idx < total) {
        int c = (int)((idx / spatial) % (IndexT)C);
        out[idx] = fused_chain(in[idx], sum_t[c]);
    }
}

template <typename IndexT>
__global__ void fused_chain_vec4_kernel(const float4* __restrict__ in,
                                        float4* __restrict__ out,
                                        const float* __restrict__ sum_t,
                                        int C, IndexT spatial4, IndexT total4) {
    IndexT idx = (IndexT)blockIdx.x * (IndexT)blockDim.x + (IndexT)threadIdx.x;
    if (idx < total4) {
        // spatial4 = D*H*W/4, so every float4 stays inside one channel
        int c = (int)((idx / spatial4) % (IndexT)C);
        float s = sum_t[c];
        float4 v = in[idx];
        v.x = fused_chain(v.x, s);
        v.y = fused_chain(v.y, s);
        v.z = fused_chain(v.z, s);
        v.w = fused_chain(v.w, s);
        out[idx] = v;
    }
}

} // namespace

torch::Tensor fused_postprocess(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");
    TORCH_CHECK(sum_t.is_cuda(), "sum_t must be a CUDA tensor");
    TORCH_CHECK(sum_t.scalar_type() == at::kFloat, "sum_t must be float32");

    const at::cuda::OptionalCUDAGuard guard(at::device_of(x));

    auto xc = x.contiguous();
    auto sc = sum_t.contiguous();
    auto out = torch::empty_like(xc);

    const int C = (int)xc.size(1);
    long long spatial = 1;
    for (int i = 2; i < xc.dim(); ++i) spatial *= (long long)xc.size(i);
    const long long total = (long long)xc.numel();
    if (total == 0) return out;

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

    const bool aligned16 =
        ((reinterpret_cast<uintptr_t>(xc.data_ptr<float>()) & 0xF) == 0) &&
        ((reinterpret_cast<uintptr_t>(out.data_ptr<float>()) & 0xF) == 0);

    if (aligned16 && (spatial % 4 == 0) && (total % 4 == 0)) {
        const long long total4 = total >> 2;
        const long long spatial4 = spatial >> 2;
        const long long nblocks = (total4 + block - 1) / block;
        if (total4 < (long long)INT_MAX - 2 * block) {
            fused_chain_vec4_kernel<int><<<(unsigned int)nblocks, block, 0, stream>>>(
                reinterpret_cast<const float4*>(xc.data_ptr<float>()),
                reinterpret_cast<float4*>(out.data_ptr<float>()),
                sc.data_ptr<float>(), C, (int)spatial4, (int)total4);
        } else {
            fused_chain_vec4_kernel<long long><<<(unsigned int)nblocks, block, 0, stream>>>(
                reinterpret_cast<const float4*>(xc.data_ptr<float>()),
                reinterpret_cast<float4*>(out.data_ptr<float>()),
                sc.data_ptr<float>(), C, (long long)spatial4, (long long)total4);
        }
    } else {
        const long long nblocks = (total + block - 1) / block;
        if (total < (long long)INT_MAX - 2 * block) {
            fused_chain_scalar_kernel<int><<<(unsigned int)nblocks, block, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), sc.data_ptr<float>(),
                C, (int)spatial, (int)total);
        } else {
            fused_chain_scalar_kernel<long long><<<(unsigned int)nblocks, block, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), sc.data_ptr<float>(),
                C, (long long)spatial, (long long)total);
        }
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_post_cpp_source = (
    "torch::Tensor fused_postprocess(torch::Tensor x, torch::Tensor sum_t);"
)

fused_post_ext = load_inline(
    name="fused_postprocess_ext",
    cpp_sources=fused_post_cpp_source,
    cuda_sources=fused_post_cuda_source,
    functions=["fused_postprocess"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: 3D convolution is delegated to cuDNN (with autotuning
    enabled), and the entire elementwise chain
    (LeakyReLU + broadcast add + clamp + GELU) is executed by a single
    fused custom CUDA kernel, eliminating 3 intermediate tensors and
    3 extra global-memory round trips.
    """
    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(ModelNew, self).__init__()
        # Let cuDNN autotune the fastest conv algorithm for the fixed shapes.
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
        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_post_ext.fused_postprocess(x, self.sum_tensor)
        return x
```