import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a batch matrix multiplication, instance normalization, summation, residual addition, and multiplication.
    """
    def __init__(self, in_features, out_features, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.bmm = nn.Linear(in_features, out_features)
        self.instance_norm = nn.InstanceNorm2d(out_features, eps=eps, momentum=momentum)

    def forward(self, x, y):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).
            y (torch.Tensor): Input tensor of shape (batch_size, out_features).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.bmm(x)
        x = self.instance_norm(x.unsqueeze(1).unsqueeze(1)).squeeze(1).squeeze(1)
        x = x + y
        x = x * y
        return x

batch_size = 1024  # Increased batch size
in_features = 8192  # Increased input features
out_features = 8192  # Increased output features

def get_inputs():
    return [torch.rand(batch_size, in_features), torch.rand(batch_size, out_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 fused_forward(torch::Tensor x, torch::Tensor y, torch::Tensor weight_t, torch::Tensor bias, double eps);
'''

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

__inline__ __device__ float warp_sum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__inline__ __device__ float block_sum(float v) {
    __shared__ float s[32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    v = warp_sum(v);
    if (lane == 0) s[wid] = v;
    __syncthreads();
    v = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? s[lane] : 0.0f;
    if (wid == 0) v = warp_sum(v);
    return v;
}

__global__ __launch_bounds__(256, 2)
void norm_bias_residual_mul_kernel(float* __restrict__ z,
                                   const float* __restrict__ y,
                                   const float* __restrict__ bias,
                                   int B, int O, float eps) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* zrow = z + (long long)row * O;
    const float* yrow = y + (long long)row * O;
    float* out = z + (long long)row * O;

    int O4 = O >> 2;
    const float4* z4 = reinterpret_cast<const float4*>(zrow);
    const float4* b4 = reinterpret_cast<const float4*>(bias);

    float sum = 0.0f;
    float sumsq = 0.0f;
    for (int i = tid; i < O4; i += blockDim.x) {
        float4 a = z4[i];
        float4 b = b4[i];
        a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w;
        sum   += a.x + a.y + a.z + a.w;
        sumsq += a.x*a.x + a.y*a.y + a.z*a.z + a.w*a.w;
    }
    for (int j = (O4 << 2) + tid; j < O; j += blockDim.x) {
        float a = zrow[j] + bias[j];
        sum += a;
        sumsq += a * a;
    }

    float total = block_sum(sum);
    float total2 = block_sum(sumsq);

    __shared__ float s_mean, s_invstd;
    if (tid == 0) {
        float mean = total / (float)O;
        float var = total2 / (float)O - mean * mean;
        if (var < 0.0f) var = 0.0f;
        s_mean = mean;
        s_invstd = rsqrtf(var + eps);
    }
    __syncthreads();

    float mean = s_mean;
    float invstd = s_invstd;
    const float4* y4 = reinterpret_cast<const float4*>(yrow);
    float4* out4 = reinterpret_cast<float4*>(out);
    for (int i = tid; i < O4; i += blockDim.x) {
        float4 a = z4[i];
        float4 yy = y4[i];
        float4 b = b4[i];
        a.x = ((a.x + b.x - mean) * invstd + yy.x) * yy.x;
        a.y = ((a.y + b.y - mean) * invstd + yy.y) * yy.y;
        a.z = ((a.z + b.z - mean) * invstd + yy.z) * yy.z;
        a.w = ((a.w + b.w - mean) * invstd + yy.w) * yy.w;
        out4[i] = a;
    }
    for (int j = (O4 << 2) + tid; j < O; j += blockDim.x) {
        float yy = yrow[j];
        out[j] = ((zrow[j] + bias[j] - mean) * invstd + yy) * yy;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor y, torch::Tensor weight_t, torch::Tensor bias, double eps) {
    auto z = at::mm(x, weight_t);
    norm_bias_residual_mul_kernel<<<z.size(0), 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        z.data_ptr<float>(), y.data_ptr<float>(), bias.data_ptr<float>(), z.size(0), z.size(1), (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return z;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.bmm = nn.Linear(in_features, out_features)
        self.eps = eps
        self.momentum = momentum
        self._weight_t = None
        self._weight_version = -1

    def forward(self, x, y):
        w = self.bmm.weight
        if self._weight_t is None or self._weight_version != w._version or self._weight_t.device != w.device:
            self._weight_t = w.detach().t().contiguous()
            self._weight_version = w._version
        return _ext.fused_forward(x.contiguous(), y.contiguous(), self._weight_t, self.bmm.bias, float(self.eps))