import math
from dataclasses import dataclass
import torch
from torch import nn
import torch.nn.functional as F
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

CPP_WRAPPER = """
void batched_compute(torch::Tensor a, torch::Tensor b, torch::Tensor c);
void permuted_compute(torch::Tensor a, torch::Tensor b, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp16.h>

using f32x16 = __attribute__( (__vector_size__(16 * sizeof(float)) )) float;
using bf16x4 = uint64_t;

__device__ static inline f32x16 __builtin_amdgcn_mfma_f32_32x32x8bf16(bf16x4 a, bf16x4 b, f32x16 d, size_t cbsz, size_t abid, size_t blgp)
{
    asm volatile("v_mfma_f32_32x32x8_bf16 %0, %1, %2, %3" : "+v"(d) : "v"(a), "v"(b), "v"(d));
    return d;
}

__device__ static uint64_t pack_f32x4_to_u64(const float* src) {
    hip_bfloat16 dst[] = { (hip_bfloat16)(src[0]), (hip_bfloat16)(src[1]), (hip_bfloat16)(src[2]), (hip_bfloat16)(src[3]) };
    return *(uint64_t*)dst;
}

__global__ void batched_compute_kernel(const uint16_t* a, const uint16_t* b, uint64_t* c,
                                       size_t m, size_t n, size_t k,
                                       size_t a0, size_t a1, size_t a2, 
                                       size_t b0, size_t b1, size_t b2) {
    const size_t cx = 32 * blockIdx.x;
    const size_t cy = 32 * blockIdx.y;
    const size_t tx = threadIdx.x & 31;
    const size_t ty = threadIdx.x >> 5 & 1;
    a += a0 * blockIdx.z + a1 * (cx + tx) + a2 * ty * 4;
    b += b0 * blockIdx.z + b1 * (cy + tx) + b2 * ty * 4;
    c += blockIdx.z * m * n >> 2;
    f32x16 result = { };
    for(size_t i = 0; i < k; i += 8) {
        uint64_t ap = a[a2 * i], bp = b[b2 * i];
        ap |= (uint64_t)a[a2 * (i + 1)] << 16;
        ap |= (uint64_t)a[a2 * (i + 2)] << 32;
        ap |= (uint64_t)a[a2 * (i + 3)] << 48;

        bp |= (uint64_t)b[b2 * (i + 1)] << 16;
        bp |= (uint64_t)b[b2 * (i + 2)] << 32;
        bp |= (uint64_t)b[b2 * (i + 3)] << 48;

        result = __builtin_amdgcn_mfma_f32_32x32x8bf16(bp, ap, result, 0, 0, 0);
    }
    const auto* src = reinterpret_cast<const float(*)[4]>(&result);
    uint64_t* dst = c + ((tx + cx) * n + cy) / 4 + ty;
    dst[0] = pack_f32x4_to_u64(src[0]);
    dst[2] = pack_f32x4_to_u64(src[1]);
    dst[4] = pack_f32x4_to_u64(src[2]);
    dst[6] = pack_f32x4_to_u64(src[3]);
}

__global__ void permuted_compute_kernel(const bf16x4* a, const bf16x4* b, uint64_t* c, size_t m, size_t n, size_t k) {
    size_t cx = 32 * blockIdx.x;
    size_t cy = 32 * blockIdx.y;
    const size_t col = k * 64;
    size_t tx = threadIdx.x & 31;
    size_t ty = threadIdx.x >> 5 & 1;
    a += blockIdx.x * col | threadIdx.x;
    b += blockIdx.y * col | threadIdx.x;
    f32x16 result = { };
    for(size_t i = 0; i < k; ++i) {
        result = __builtin_amdgcn_mfma_f32_32x32x8bf16(b[i * 64], a[i * 64], result, 0, 0, 0);
    }
    const auto* src = reinterpret_cast<const float(*)[4]>(&result);
    uint64_t* dst = c + ((tx + cx) * n + cy) / 4 + ty;
    dst[0] = pack_f32x4_to_u64(src[0]);
    dst[2] = pack_f32x4_to_u64(src[1]);
    dst[4] = pack_f32x4_to_u64(src[2]);
    dst[6] = pack_f32x4_to_u64(src[3]);
}

void batched_compute(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
    size_t count = a.size(0);
    size_t m = a.size(1);
    size_t n = b.size(1);
    size_t k = a.size(2);
    batched_compute_kernel<<<dim3(m/32, n/32, count), dim3(64), 0, 0>>> (
        (uint16_t*)a.data_ptr(), (uint16_t*)b.data_ptr(), (uint64_t*)c.data_ptr(),
        m, n, k,
        (size_t)a.stride(0), (size_t)a.stride(1), (size_t)a.stride(2),
        (size_t)b.stride(0), (size_t)b.stride(1), (size_t)b.stride(2));
}

void permuted_compute(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
    size_t n = c.size(-1);
    size_t m = c.numel() / n;
    size_t k = a.size(-2);
    permuted_compute_kernel<<<dim3(m/32, n/32), dim3(64), 0, 0>>> ((bf16x4*)a.data_ptr(), (bf16x4*)b.data_ptr(), (uint64_t*)c.data_ptr(), m, n, k);
}

"""

import os
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
os.environ["CXX"] = "clang++"

module = load_inline(
    name='compute',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['batched_compute', 'permuted_compute'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)

class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        theta = 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (d_model//2))
        self.register_buffer("theta", theta)

    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(-2)
        d_model = x.size(-1)
        assert d_model == self.d_model
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        return x * cos + self.rotate_half(x) * sin


class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        self.register_buffer('data', torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device='cuda'))
        self.seq_len = 0
        self.zero()

    def zero(self) -> None:
        self.data.zero_()
    
    def get_data(self) -> torch.Tensor:
        return self.data

    def forward(self, c_kv: torch.Tensor) -> torch.Tensor:
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"

        self.data = self.data.to(c_kv.dtype)
        self.data[
            :, self.seq_len : self.seq_len + c_kv.size(1), :
        ] = c_kv
        self.seq_len += c_kv.size(1)

        return self.data[:, :self.seq_len], self.seq_len

PAD_N = 32
PAD_K = 8

def permuteMatrix(mat: torch.Tensor) -> torch.Tensor:
    assert 2 == mat.dim()
    k = mat.size(-1)
    n = mat.size(-2)
    assert k % 8  == 0
    assert n % 32 == 0
    k = k // 8
    n = n // 32
    return mat.view(n, 32, k, 2, 4).permute(0, 2, 3, 1, 4).reshape(n, k, 256)


def matmul(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    weight = permuteMatrix(weight)
    size = list(input.size())
    input = input.reshape(input.numel() // size[-1], size[-1])
    size[-1] = weight.size(-3) * 32
    m = input.size(-2)
    pad = (-m) % PAD_N
    if 0 != pad:
        input = F.pad(input, (0, pad, 0, 0), mode='constant', value=0)
    input = permuteMatrix(input)
    assert input.size(-2) == weight.size(-2)
    output = torch.empty((m + pad, size[-1]), dtype=torch.bfloat16, device='cuda')
    module.permuted_compute(input, weight, output)
    if 0 != pad:
        output = output[-pad, :]
    return output.view(size)

def padmat(a: torch.Tensor) -> torch.Tensor:
    n = a.size(1)
    k = a.size(2)
    npad = (-n) % PAD_N
    kpad = (-k) % PAD_K
    if 0 == npad and 0 == kpad:
        return a
    return F.pad(a, (0, kpad, 0, npad, 0, 0), mode='constant', value=0)

# def batchmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
#     a1 = padmat(a)
#     b1 = padmat(b)
#     mpad = a1.size(-2) - a.size(-2)
#     npad = b1.size(-2) - b.size(-2)
#     count = a1.size(0)
#     m = a1.size(-2) // 32
#     n = b1.size(-2) // 32
#     k = a1.size(-1) // 8
#     input = torch.cat((a1, b1), dim=1).view(count, m + n, 32, k, 2, 4).permute(0, 1, 3, 4, 2, 5).reshape(count, m + n, k, 256)
#     c = torch.empty((count, a1.size(-2), b1.size(-2)), dtype=a.dtype, device=a.device)
#     module.batch_compute(input, c)
#     if 0 != mpad and 0 != npad:
#         return c[:, :-mpad, :-npad]
#     if 0 != mpad:
#         return c[:, :-mpad, :]
#     if 0 != npad:
#         return c[:, :, :-npad]
#     return c



def bmatmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    n = b.size(-2)
    input = a.flatten(start_dim=1, end_dim=2)
    output = torch.bmm(input, b.transpose(-1, -2))
    return output.view(a.size(0), a.size(1), a.size(2), n)

def bmatmul1(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    n = b.size(-2)
    input = a.flatten(start_dim=1, end_dim=2)
    m = input.size(-2)
    count = input.size(0)

    output = torch.empty((count, m, n), dtype=torch.bfloat16, device='cuda')
    module.batched_compute(input, b, output)
    return output.view(a.size(0), a.size(1), a.size(2), n)
    
@dataclass
class Config:
    batch_size: int
    dim: int
    n_heads: int
    q_lora_rank: int 
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int
    seq_len: int
    max_seq_len: int
    kv_cache_shape: tuple
    Q_proj_down_weight: torch.Tensor
    Q_proj_up_weight: torch.Tensor
    KV_proj_down_weight: torch.Tensor
    KV_proj_up_weight: torch.Tensor
    wo_weight: torch.Tensor

class MLA(nn.Module):
    def __init__(self, config: Config, seq_len: int):
        super().__init__()
        self.dim = config.dim
        self.n_heads = config.n_heads
        self.q_lora_rank = config.q_lora_rank
        self.kv_lora_rank = config.kv_lora_rank
        self.nope_head_dim = config.qk_nope_head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.v_head_dim = config.v_head_dim
        # Down-projection matrices
        self.Q_proj_down_weight = config.Q_proj_down_weight
        self.Q_proj_up_weight = config.Q_proj_up_weight
        self.KV_proj_down_weight = config.KV_proj_down_weight
        self.KV_proj_up_weight = config.KV_proj_up_weight
        self.wo_weight = config.wo_weight

        # RoPE on half embeddings
        self.q_rope = RoPE(self.rope_head_dim)
        self.k_rope = RoPE(self.rope_head_dim)

        # Output projection
        self.eps = 1e-6
   
    def forward(self, x: torch.Tensor, kv_cache: KVCache) -> torch.Tensor:
        # seq_len = 1 always here
        batch_size, seq_len, model_dim = x.size()

        ################################################################################
        #                 Step 1: Handle down-projection + KV cache                    #
        ################################################################################
        q_lora = matmul(x, self.Q_proj_down_weight)
        kv_lora = matmul(x, self.KV_proj_down_weight)
        kv_lora, kv_len = kv_cache(kv_lora)
        query_pos = kv_len - 1

        ################################################################################
        #                  Step 2: Up-project and prepare NoPE + RoPE                  #
        ################################################################################

        # Handle queries Q first
        q_nope_and_rope = matmul(q_lora, self.Q_proj_up_weight).view(
            batch_size, seq_len, self.n_heads, self.nope_head_dim + self.rope_head_dim)
        q_nope, q_rope = torch.split(q_nope_and_rope, [self.nope_head_dim, self.rope_head_dim], dim=-1)

        # Compute RoPE for queries and combine with no-RoPE part
        q_rope = q_rope.permute(2, 0, 1, 3) # n_heads x bs x seq_len x rope_head_dim
        q_rope = self.q_rope(q_rope, start_pos=query_pos)

        q_nope = q_nope.permute(2, 0, 1, 3) # n_heads x bs x seq_len x rope_head_dim

        # Handle keys and values K/V. V does not need RoPE
        kv_nope, k_rope = torch.split(kv_lora, [self.kv_lora_rank, self.rope_head_dim], dim=-1)
        # Compute RoPE for keys and combine with no-RoPE part
        k_rope = self.k_rope(k_rope)
        kv_lora = torch.cat((kv_nope, k_rope), dim=-1)
                
        ################################################################################
        #                        Compute Multi-head Attention                          #
        ################################################################################

        KV_proj_up_weight = self.KV_proj_up_weight.view(self.n_heads, self.nope_head_dim + self.v_head_dim, -1)
        K_proj_up_weight, V_proj_up_weight = KV_proj_up_weight.split((self.nope_head_dim, self.v_head_dim), 1)
        K_proj_up_weight = K_proj_up_weight.transpose(-1, -2)
        kv_nope = kv_nope.transpose(-1, -2)
        scores = bmatmul(q_nope, K_proj_up_weight)
        scores = torch.cat((scores, q_rope), dim=-1)
        scores = bmatmul(scores.transpose(0, 1), kv_lora) * (1.0 / math.sqrt(self.rope_head_dim + self.nope_head_dim))
        attn = F.softmax(scores, dim=-1).to(torch.bfloat16)
        y = bmatmul(attn, kv_nope)
        y = bmatmul(y.transpose(0, 1), V_proj_up_weight)
        y = y.transpose(0, 1).reshape(batch_size, 1, -1)
        y = matmul(y, self.wo_weight)

        return y, kv_cache.get_data()

def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    model = MLA(config, kv_cache.seq_len).to('cuda')
    output, kv_cache = model(x, kv_cache)
    return output, kv_cache
python eval.py benchmark /tmp/tmp_ks_1tos
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/compute/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/compute/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/compute/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/compute/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 2
[1/1] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o compute.so
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
