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

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

# Fused CUDA kernel:
#   y = GELU(clamp(LeakyReLU(conv_out, 0.2) + sum_tensor[c], -1, 1))
# The sum_tensor is per-channel (C,1,1,1), so each thread computes its
# channel index from the linear index. Vectorized float4 path is used when
# the per-channel spatial size (D*H*W) is divisible by 4 (guaranteeing a
# float4 never straddles two channels) and pointers are 16B aligned.
fused_act_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float gelu_exact(float x) {
    // exact (erf-based) GELU, matches F.gelu default
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f));
}

__device__ __forceinline__ float fused_op(float v, float s) {
    v = (v > 0.0f) ? v : v * 0.2f;          // LeakyReLU(negative_slope=0.2)
    v = v + s;                              // add per-channel sum_tensor
    v = fminf(fmaxf(v, -1.0f), 1.0f);       // clamp(-1, 1)
    return gelu_exact(v);                   // GELU
}

template <typename index_t>
__global__ void fused_act_vec4_kernel(const float4* __restrict__ x,
                                      const float* __restrict__ sum_t,
                                      float4* __restrict__ out,
                                      index_t num_vec4,
                                      int C,
                                      index_t dhw) {
    index_t i = (index_t)blockIdx.x * blockDim.x + threadIdx.x;
    index_t stride = (index_t)gridDim.x * blockDim.x;
    for (; i < num_vec4; i += stride) {
        index_t base = i * 4;
        int c = (int)((base / dhw) % (index_t)C);
        float s = __ldg(&sum_t[c]);
        float4 v = x[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);
        out[i] = r;
    }
}

template <typename index_t>
__global__ void fused_act_scalar_kernel(const float* __restrict__ x,
                                        const float* __restrict__ sum_t,
                                        float* __restrict__ out,
                                        index_t n,
                                        int C,
                                        index_t dhw) {
    index_t i = (index_t)blockIdx.x * blockDim.x + threadIdx.x;
    index_t stride = (index_t)gridDim.x * blockDim.x;
    for (; i < n; i += stride) {
        int c = (int)((i / dhw) % (index_t)C);
        out[i] = fused_op(x[i], __ldg(&sum_t[c]));
    }
}

torch::Tensor fused_act_forward(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");
    TORCH_CHECK(sum_t.scalar_type() == at::kFloat, "sum_t must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(sum_t.is_contiguous(), "sum_t must be contiguous");

    auto out = torch::empty_like(x);
    const long long n = x.numel();
    if (n == 0) return out;

    const int C = (int)x.size(1);
    const long long dhw = (long long)x.size(2) * x.size(3) * x.size(4);

    const int block = 256;
    const long long max_blocks = 2097152LL;  // plenty; grid-stride covers rest
    auto stream = at::cuda::getCurrentCUDAStream();

    const bool vec_ok = (dhw % 4 == 0) && (n % 4 == 0) &&
                        ((reinterpret_cast<uintptr_t>(x.data_ptr<float>()) & 15) == 0) &&
                        ((reinterpret_cast<uintptr_t>(out.data_ptr<float>()) & 15) == 0);

    if (vec_ok) {
        const long long num_vec4 = n / 4;
        long long blocks = (num_vec4 + block - 1) / block;
        if (blocks > max_blocks) blocks = max_blocks;
        if (num_vec4 <= 0x7fffffffLL) {
            fused_act_vec4_kernel<int><<<(int)blocks, block, 0, stream>>>(
                reinterpret_cast<const float4*>(x.data_ptr<float>()),
                sum_t.data_ptr<float>(),
                reinterpret_cast<float4*>(out.data_ptr<float>()),
                (int)num_vec4, C, (int)dhw);
        } else {
            fused_act_vec4_kernel<long long><<<(int)blocks, block, 0, stream>>>(
                reinterpret_cast<const float4*>(x.data_ptr<float>()),
                sum_t.data_ptr<float>(),
                reinterpret_cast<float4*>(out.data_ptr<float>()),
                num_vec4, C, dhw);
        }
    } else {
        long long blocks = (n + block - 1) / block;
        if (blocks > max_blocks) blocks = max_blocks;
        if (n <= 0x7fffffffLL) {
            fused_act_scalar_kernel<int><<<(int)blocks, block, 0, stream>>>(
                x.data_ptr<float>(), sum_t.data_ptr<float>(),
                out.data_ptr<float>(), (int)n, C, (int)dhw);
        } else {
            fused_act_scalar_kernel<long long><<<(int)blocks, block, 0, stream>>>(
                x.data_ptr<float>(), sum_t.data_ptr<float>(),
                out.data_ptr<float>(), n, C, dhw);
        }
    }
    return out;
}
"""

fused_act_cpp_source = (
    "torch::Tensor fused_act_forward(torch::Tensor x, torch::Tensor sum_t);"
)

# Compile the fused activation kernel
fused_act = load_inline(
    name="fused_leakyrelu_add_clamp_gelu",
    cpp_sources=fused_act_cpp_source,
    cuda_sources=fused_act_source,
    functions=["fused_act_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - Conv3d is kept (cuDNN, benchmark-tuned).
      - LeakyReLU + per-channel tensor add + clamp + GELU are fused into a
        single memory-bound custom CUDA kernel (float4 vectorized), removing
        three full tensor read/write passes and three kernel launches.
    """
    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_act.fused_act_forward(x, self.sum_tensor)
        return x
```