```
# 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: (x - subtract_value_1 - subtract_value_2) -> Mish
# into a single memory-bound pass over the convolution output.
#
# Mish identity used for speed and stability:
#   mish(z) = z * tanh(softplus(z))
#   tanh(softplus(z)) = ((1 + e^z)^2 - 1) / ((1 + e^z)^2 + 1)
# so only one expf plus a few FMAs / one divide are needed per element.
fused_sub_mish_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float sub_mish_op(float z) {
    // For z > 20, tanh(softplus(z)) rounds to exactly 1.0f in fp32
    // (this matches ATen's softplus threshold), so mish(z) == z.
    if (z > 20.0f) {
        return z;
    }
    float t = expf(z);
    float u = t * (t + 2.0f);   // (1 + e^z)^2 - 1
    return z * (u / (u + 2.0f));
}

__global__ void sub_mish_vec4_kernel(float4* __restrict__ data, int n4, float c) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = data[i];
        v.x = sub_mish_op(v.x - c);
        v.y = sub_mish_op(v.y - c);
        v.z = sub_mish_op(v.z - c);
        v.w = sub_mish_op(v.w - c);
        data[i] = v;
    }
}

__global__ void sub_mish_tail_kernel(float* __restrict__ data,
                                     long long start, long long n, float c) {
    long long i = start + (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        data[i] = sub_mish_op(data[i] - c);
    }
}

torch::Tensor sub_mish_cuda(torch::Tensor x, double c) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");

    auto out = x.contiguous();  // same tensor if already contiguous
    const long long n = out.numel();
    if (n == 0) return out;

    float* ptr = out.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int block = 256;
    const long long n4 = n >> 2;
    if (n4 > 0) {
        const int grid = (int)((n4 + block - 1) / block);
        sub_mish_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<float4*>(ptr), (int)n4, (float)c);
    }
    const long long rem_start = n4 << 2;
    if (rem_start < n) {
        const int rem = (int)(n - rem_start);
        sub_mish_tail_kernel<<<(rem + block - 1) / block, block, 0, stream>>>(
            ptr, rem_start, n, (float)c);
    }
    return out;
}
"""

fused_sub_mish_cpp_source = (
    "torch::Tensor sub_mish_cuda(torch::Tensor x, double c);"
)

fused_sub_mish = load_inline(
    name="fused_sub_mish_ext",
    cpp_sources=fused_sub_mish_cpp_source,
    cuda_sources=fused_sub_mish_cuda_source,
    functions=["sub_mish_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN convolution followed by a single fused custom CUDA
    kernel that performs both subtractions and the Mish activation in one pass.
    """
    def __init__(self, in_channels, out_channels, kernel_size, subtract_value_1, subtract_value_2):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        # The two scalar subtractions are folded into a single constant.
        self.sub_total = float(subtract_value_1) + float(subtract_value_2)
        # Let cuDNN autotune the fastest convolution algorithm for this shape.
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        x = self.conv(x)
        x = fused_sub_mish.sub_mish_cuda(x, self.sub_total)
        return x
```