import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, embed_dim, num_heads):
        """
        Attention Block using Multihead Self-Attention.
        :param embed_dim: Embedding dimension (the number of channels)
        :param num_heads: Number of attention heads
        """
        super(Model, self).__init__()
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        """
        Forward pass of the AttentionBlock.
        :param x: Input tensor of shape (B, C, H, W)
        :return: Output tensor of the same shape (B, C, H, W)
        """
        B, C, H, W = x.shape
        x = x.view(B, C, H * W).permute(2, 0, 1)  # (seq_len, batch_size, embed_dim)
        attn_output, _ = self.attn(x, x, x)
        x = self.norm(attn_output + x)  # (seq_len, batch_size, embed_dim)
        x = x.permute(1, 2, 0).view(B, C, H, W)
        return x

embed_dim = 128
num_heads = 4
batch_size = 2
num_channels = embed_dim
image_height = 128
image_width = 128

def get_inputs():
    return [torch.rand(batch_size, num_channels, image_height, image_width)]

def get_init_inputs():
    return [embed_dim, num_heads]

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_ln_out(torch::Tensor proj, torch::Tensor x_blc, torch::Tensor ln_w, torch::Tensor ln_b, int64_t H, int64_t W);

torch::Tensor attn_forward(torch::Tensor x,
                           torch::Tensor in_w, torch::Tensor in_b,
                           torch::Tensor out_w, torch::Tensor out_b,
                           torch::Tensor ln_w, torch::Tensor ln_b,
                           int64_t num_heads) {
    auto B = x.size(0);
    auto C = x.size(1);
    auto H = x.size(2);
    auto W = x.size(3);
    auto L = H * W;
    auto D = C / num_heads;

    auto x_blc = x.view({B, C, L}).permute({0, 2, 1}).contiguous();

    auto qkv = at::linear(x_blc, in_w, in_b);
    qkv = qkv.view({B, L, 3, num_heads, D}).permute({2, 0, 3, 1, 4});
    auto q = qkv.select(0, 0);
    auto k = qkv.select(0, 1);
    auto v = qkv.select(0, 2);

    auto attn = at::scaled_dot_product_attention(q, k, v, c10::nullopt, 0.0, false, c10::nullopt);
    auto attn_blc = attn.permute({0, 2, 1, 3}).contiguous().view({B, L, C});
    auto proj = at::linear(attn_blc, out_w, out_b);

    return fused_ln_out(proj, x_blc, ln_w, ln_b, H, W);
}
'''

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

__global__ void ln_out_kernel(const float* __restrict__ proj,
                              const float* __restrict__ x,
                              const float* __restrict__ gamma,
                              const float* __restrict__ beta,
                              float* __restrict__ out,
                              int B, int L, int C) {
    int t = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float smem[];

    float v = 0.0f;
    if (tid < C) {
        int idx = t * C + tid;
        v = proj[idx] + x[idx];
    }

    smem[tid] = (tid < C) ? v : 0.0f;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    float mean = smem[0] / (float)C;
    __syncthreads();

    float diff = (tid < C) ? (v - mean) : 0.0f;
    smem[tid] = diff * diff;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    float inv_std = rsqrtf(smem[0] / (float)C + 1.0e-5f);

    if (tid < C) {
        int b = t / L;
        int l = t - b * L;
        int out_idx = (b * C + tid) * L + l;
        out[out_idx] = diff * inv_std * gamma[tid] + beta[tid];
    }
}

torch::Tensor fused_ln_out(torch::Tensor proj,
                           torch::Tensor x_blc,
                           torch::Tensor ln_w,
                           torch::Tensor ln_b,
                           int64_t H,
                           int64_t W) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(proj));
    int B = (int)proj.size(0);
    int L = (int)proj.size(1);
    int C = (int)proj.size(2);

    auto out = torch::empty({B, C, H, W}, proj.options());

    int threads = 1;
    while (threads < C) threads <<= 1;
    if (threads < 32) threads = 32;

    ln_out_kernel<<<B * L, threads, threads * sizeof(float), at::cuda::getCurrentCUDAStream()>>>(
        proj.data_ptr<float>(),
        x_blc.data_ptr<float>(),
        ln_w.data_ptr<float>(),
        ln_b.data_ptr<float>(),
        out.data_ptr<float>(),
        B, L, C);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(ModelNew, self).__init__()
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.norm = nn.LayerNorm(embed_dim)
        self.num_heads = num_heads

    def forward(self, x):
        return _ext.attn_forward(
            x,
            self.attn.in_proj_weight,
            self.attn.in_proj_bias,
            self.attn.out_proj.weight,
            self.attn.out_proj.bias,
            self.norm.weight,
            self.norm.bias,
            self.num_heads,
        )