import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies Mish, and applies Mish again.
    """
    def __init__(self, in_features, out_features):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        x = self.linear(x)
        x = torch.nn.functional.mish(x)
        x = torch.nn.functional.mish(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192

def get_inputs():
    return [torch.rand(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor linear_mish2(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>

__device__ __forceinline__ float sp_f(float x) { return (x > 20.0f) ? x : log1pf(expf(x)); }
__device__ __forceinline__ float mish_f(float x) { return x * tanhf(sp_f(x)); }
__device__ __forceinline__ float mish2_bias(float x, float b) { float z=x+b; float a=mish_f(z); return mish_f(a); }

__global__ void fused_bias_mish2_vec4(float* __restrict__ y, const float* __restrict__ bias, long long n4) {
    long long i4 = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; i4 < n4; i4 += stride) {
        long long elem = i4 << 2;
        int col = (int)elem & 8191;
        float4 v = *reinterpret_cast<float4*>(y + elem);
        float4 b = *reinterpret_cast<const float4*>(bias + col);
        v.x = mish2_bias(v.x, b.x);
        v.y = mish2_bias(v.y, b.y);
        v.z = mish2_bias(v.z, b.z);
        v.w = mish2_bias(v.w, b.w);
        *reinterpret_cast<float4*>(y + elem) = v;
    }
}

__global__ void fused_bias_mish2_scalar(float* __restrict__ y, const float* __restrict__ bias, long long n, int out_features) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; idx < n; idx += stride) {
        int col = (out_features == 8192) ? ((int)idx & 8191) : (int)(idx % out_features);
        y[idx] = mish2_bias(y[idx], bias[col]);
    }
}

torch::Tensor linear_mish2(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    auto y = torch::empty({x.size(0), weight.size(0)}, x.options());
    auto wt = weight.t();
    torch::mm_out(y, x, wt);
    const long long n = y.numel();
    if (n) {
        const int out_features = (int)y.size(1);
        if (out_features == 8192 && (n % 4 == 0)) {
            const int threads = 256; long long n4 = n >> 2;
            int blocks = (int)((n4 + threads - 1) / threads); if (blocks > 65535) blocks = 65535;
            fused_bias_mish2_vec4<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), bias.data_ptr<float>(), n4);
        } else {
            const int threads = 256; int blocks = (int)((n + threads - 1) / threads); if (blocks > 65535) blocks = 65535;
            fused_bias_mish2_scalar<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), bias.data_ptr<float>(), n, out_features);
        }
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return y;
}
'''
_ext = load_inline(name='linear_mish2_ext_vec4_final', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['linear_mish2'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__(); self.linear=nn.Linear(in_features,out_features)
    def forward(self,x):
        return _ext.linear_mish2(x, self.linear.weight, self.linear.bias)