import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a matrix multiplication, batch normalization, bias addition, division, and Swish activation.
    """
    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, bias_shape=(1,), divide_value=1.0):
        super(Model, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(out_features, eps=bn_eps, momentum=bn_momentum)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.divide_value = divide_value

    def forward(self, x):
        x = self.matmul(x)
        x = self.bn(x)
        x = x + self.bias
        x = x / self.divide_value
        x = x * torch.sigmoid(x)
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
bn_eps = 1e-5
bn_momentum = 0.1
bias_shape = (1,)
divide_value = 1.0

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

def get_init_inputs():
    return [in_features, out_features, bn_eps, bn_momentum, bias_shape, divide_value]

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


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

torch::Tensor fused_forward(torch::Tensor input,
                           torch::Tensor weight,
                           torch::Tensor lin_bias,
                           torch::Tensor bn_weight,
                           torch::Tensor bn_bias,
                           torch::Tensor running_mean,
                           torch::Tensor running_var,
                           torch::Tensor extra_bias,
                           bool training,
                           double eps,
                           double momentum,
                           double divide_value);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/ops/linear.h>
#include <ATen/ops/batch_norm.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void bias_swish_kernel(const float* __restrict__ in,
                                  const float* __restrict__ bias,
                                  float* __restrict__ out,
                                  long long total,
                                  int C,
                                  long long bias_numel,
                                  float inv_div) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    float b;
    if (bias_numel == 1) {
        b = bias[0];
    } else if (bias_numel == (long long)C) {
        b = bias[idx % C];
    } else if (bias_numel == total) {
        b = bias[idx];
    } else {
        // Covers the only broadcast patterns used by this benchmark (scalar) and
        // common (C,) / full-size cases.  For unusual shapes, fall back to a
        // cyclic index rather than reading out of bounds.
        b = bias[idx % bias_numel];
    }

    float x = (in[idx] + b) * inv_div;
    float s = 1.0f / (1.0f + expf(-x));
    out[idx] = x * s;
}

torch::Tensor fused_forward(torch::Tensor input,
                           torch::Tensor weight,
                           torch::Tensor lin_bias,
                           torch::Tensor bn_weight,
                           torch::Tensor bn_bias,
                           torch::Tensor running_mean,
                           torch::Tensor running_var,
                           torch::Tensor extra_bias,
                           bool training,
                           double eps,
                           double momentum,
                           double divide_value) {
    // Keep the numerically sensitive / highly optimized GEMM and BatchNorm in
    // ATen/cuDNN, and fuse the remaining bias/divide/Swish chain into one CUDA
    // pass.  This removes several full-tensor PyTorch kernels from the forward.
    auto y = at::linear(input, weight, lin_bias);
    auto z = at::batch_norm(y, bn_weight, bn_bias, running_mean, running_var,
                            training, momentum, eps, true);

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

    const int C = (z.dim() >= 2) ? (int)z.size(1) : (int)total;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    float inv_div = 1.0f / static_cast<float>(divide_value);

    bias_swish_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        z.data_ptr<float>(), extra_bias.data_ptr<float>(), out.data_ptr<float>(),
        total, C, extra_bias.numel(), inv_div);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name="bn_linear_swish_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Drop-in replacement for the reference module.  Parameters and buffers are
    named identically so state_dict copying works unchanged.
    """
    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, bias_shape=(1,), divide_value=1.0):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(out_features, eps=bn_eps, momentum=bn_momentum)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.divide_value = divide_value

    def forward(self, x):
        return _ext.fused_forward(
            x,
            self.matmul.weight,
            self.matmul.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.bias,
            self.bn.training,
            float(self.bn.eps),
            float(self.bn.momentum if self.bn.momentum is not None else 0.1),
            float(self.divide_value),
        )