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 conv_mish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double sub1, double sub2);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float mish_f(float x) {
    if (x > 20.0f) return x;
    float e = __expf(x);
    float n = e * (e + 2.0f);
    return x * n / (n + 2.0f);
}

__global__ void conv3x3_mish_vec4_kernel(const float* __restrict__ inp,
    const float* __restrict__ w, const float* __restrict__ b, float* __restrict__ out,
    int N, int H, int W, int Co, float sub)
{
    int Ho = H - 2;
    int Wo = W - 2;
    int vecW = (Wo + 3) >> 2;
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Co * Ho * vecW;
    if (tid >= total) return;
    int vx = tid % vecW;
    long long t = tid / vecW;
    int oy = t % Ho; t /= Ho;
    int co = t % Co;
    int n = t / Co;
    int ox = vx << 2;
    float biasv = b[co] - sub;
    const float* base_in = inp + ((long long)n * 8 * H * W) + oy * W + ox;
    const float* base_w = w + (long long)co * 8 * 9;
    float* op = out + (((long long)n * Co + co) * Ho + oy) * Wo + ox;

    if (ox + 3 < Wo) {
        float acc0=biasv, acc1=biasv, acc2=biasv, acc3=biasv;
        #pragma unroll
        for (int ci=0; ci<8; ++ci) {
            const float* in_c = base_in + (long long)ci * H * W;
            const float* w_c = base_w + ci * 9;
            #pragma unroll
            for (int ky=0; ky<3; ++ky) {
                const float* ip = in_c + ky * W;
                float i0=ip[0], i1=ip[1], i2=ip[2], i3=ip[3], i4=ip[4], i5=ip[5];
                float w0=w_c[ky*3], w1=w_c[ky*3+1], w2=w_c[ky*3+2];
                acc0=fmaf(i0,w0,acc0); acc0=fmaf(i1,w1,acc0); acc0=fmaf(i2,w2,acc0);
                acc1=fmaf(i1,w0,acc1); acc1=fmaf(i2,w1,acc1); acc1=fmaf(i3,w2,acc1);
                acc2=fmaf(i2,w0,acc2); acc2=fmaf(i3,w1,acc2); acc2=fmaf(i4,w2,acc2);
                acc3=fmaf(i3,w0,acc3); acc3=fmaf(i4,w1,acc3); acc3=fmaf(i5,w2,acc3);
            }
        }
        op[0]=mish_f(acc0); op[1]=mish_f(acc1); op[2]=mish_f(acc2); op[3]=mish_f(acc3);
    } else {
        for (int j=0; j<4 && ox+j<Wo; ++j) {
            float acc = biasv;
            #pragma unroll
            for (int ci=0; ci<8; ++ci) {
                const float* in_c = inp + ((long long)n*8 + ci)*H*W + oy*W + (ox+j);
                const float* w_c = base_w + ci*9;
                #pragma unroll
                for (int ky=0; ky<3; ++ky) {
                    const float* ip = in_c + ky*W;
                    acc=fmaf(ip[0], w_c[ky*3], acc);
                    acc=fmaf(ip[1], w_c[ky*3+1], acc);
                    acc=fmaf(ip[2], w_c[ky*3+2], acc);
                }
            }
            op[j] = mish_f(acc);
        }
    }
}

torch::Tensor conv_mish_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double sub1, double sub2) {
    int N=(int)x.size(0), H=(int)x.size(2), W=(int)x.size(3), Co=(int)weight.size(0);
    int Ho=H-2, Wo=W-2;
    auto out = torch::empty({N, Co, Ho, Wo}, x.options());
    int vecW=(Wo+3)>>2;
    long long total=(long long)N*Co*Ho*vecW;
    int threads=256;
    int blocks=(int)((total+threads-1)/threads);
    conv3x3_mish_vec4_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), N,H,W,Co,(float)(sub1+sub2));
    return out;
}
'''

_ext = load_inline(name='conv_mish_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['conv_mish_forward'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    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
    def forward(self, x):
        return _ext.conv_mish_forward(x, self.conv.weight, self.conv.bias, float(self.subtract_value_1), float(self.subtract_value_2))