import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, subtracts two values, applies Mish activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, subtract_value_1, subtract_value_2):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.subtract_value_1 = subtract_value_1
        self.subtract_value_2 = subtract_value_2

    def forward(self, x):
        x = self.conv(x)
        x = x - self.subtract_value_1
        x = x - self.subtract_value_2
        x = torch.nn.functional.mish(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 256, 256
kernel_size = 3
subtract_value_1 = 0.5
subtract_value_2 = 0.2

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, subtract_value_1, subtract_value_2]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor sub_mish_inplace(torch::Tensor x, double sub);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float mish_after_sub(float a, float sub) {
    float y = a - sub;
    // Match torch softplus threshold behavior closely: softplus(y) ~= y for y > 20.
    float sp = (y > 20.0f) ? y : log1pf(expf(y));
    return y * tanhf(sp);
}

__global__ __launch_bounds__(256) void sub_mish_vec4_kernel(float4* __restrict__ data4,
                                                            long long n4,
                                                            float sub) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; i < n4; i += stride) {
        float4 v = data4[i];
        v.x = mish_after_sub(v.x, sub);
        v.y = mish_after_sub(v.y, sub);
        v.z = mish_after_sub(v.z, sub);
        v.w = mish_after_sub(v.w, sub);
        data4[i] = v;
    }
}

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

torch::Tensor sub_mish_inplace(torch::Tensor x, double sub) {
    const auto n = (long long)x.numel();
    if (n == 0) return x;
    const c10::cuda::CUDAGuard device_guard(x.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    float* ptr = x.data_ptr<float>();
    const float s = static_cast<float>(sub);
    const int threads = 256;
    const long long n4 = n >> 2;
    if (n4 > 0) {
        int blocks = (int)((n4 + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        sub_mish_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(ptr), n4, s);
    }
    const long long tail_start = n4 << 2;
    if (tail_start < n) {
        int blocks = (int)((n - tail_start + threads - 1) / threads);
        sub_mish_tail_kernel<<<blocks, threads, 0, stream>>>(ptr, tail_start, n, s);
    }
    return x;
}
"""

_ext = load_inline(
    name="sub_mish_inline_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["sub_mish_inplace"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Convolution followed by the two scalar subtractions and Mish.  The expensive
    convolution remains the same cuDNN-backed Conv2d module so state_dict layout is
    identical to the reference; the three eager pointwise passes are fused into one
    custom CUDA kernel running in-place on the convolution result.
    """
    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)
        self.subtract_value_1 = subtract_value_1
        self.subtract_value_2 = subtract_value_2
        self._sub_total = float(subtract_value_1) + float(subtract_value_2)

    def forward(self, x):
        y = self.conv(x)
        return _ext.sub_mish_inplace(y, self._sub_total)