import torch
from task import input_t, output_t
from utils import make_match_reference
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

from torch.utils.cpp_extension import load_inline

import triton
import triton.language as tl
from typing import Any, Dict, List, Optional, Tuple


block_shape = (128, 128)

_is_hip = True

@triton.autotune(configs=[
    triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':128, 'GROUP_SIZE_M': 16}, num_warps=4, num_stages=2),
    triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':128, 'GROUP_SIZE_M': 32}, num_warps=4, num_stages=2),
    triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':128, 'GROUP_SIZE_M': 16}, num_warps=4, num_stages=2),
    triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':128, 'GROUP_SIZE_M': 32}, num_warps=4, num_stages=2),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':64, 'GROUP_SIZE_M': 16}, num_warps=4, num_stages=2),
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':64, 'GROUP_SIZE_M': 32}, num_warps=4, num_stages=2),
  ],
  key=['M', 'N', 'K'],
  use_cuda_graph=True,
)
@triton.jit
def _w8a8_block_fp8_matmul_unrolledx4(
    # Pointers to inputs and output
    A,
    B,
    C,
    As,
    Bs,
    # Shape for matmul
    M,
    N,
    K,
    # Block size for block-wise quantization
    group_n,
    group_k,
    # Stride for inputs and output
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_As_m,
    stride_As_k,
    stride_Bs_k,
    stride_Bs_n,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """Triton-accelerated function used to perform linear operations (dot
    product) on input tensors `A` and `B` with block-wise quantization, and store the result in output
    tensor `C`.
    """
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    As_ptrs = As + offs_am * stride_As_m
    offs_bsn = offs_bn // group_n
    Bs_ptrs = Bs + offs_bsn * stride_Bs_n

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    # manually unroll to 4 iterations
    UNROLL_FACTOR = 1
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * UNROLL_FACTOR)):
        # 1st iteration
        # a = tl.load(
        #     a_ptrs,
        #     mask=offs_k[None, :] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
        #     other=0.0,
        # )
        # b = tl.load(
        #     b_ptrs,
        #     mask=offs_k[:, None] < K - (k * UNROLL_FACTOR) * BLOCK_SIZE_K,
        #     other=0.0,
        # )

        a = tl.load(
            a_ptrs,
        )
        b = tl.load(
            b_ptrs,
        )

        k_start = (k * UNROLL_FACTOR) * BLOCK_SIZE_K
        offs_ks = k_start // group_k
        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

        # # 2nd iteration
        # a = tl.load(
        #     a_ptrs,
        #     mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
        #     other=0.0,
        # )
        # b = tl.load(
        #     b_ptrs,
        #     mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 1) * BLOCK_SIZE_K,
        #     other=0.0,
        # )

        # k_start = k_start + BLOCK_SIZE_K
        # offs_ks = k_start // group_k
        # a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
        # b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

        # accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        # a_ptrs += BLOCK_SIZE_K * stride_ak
        # b_ptrs += BLOCK_SIZE_K * stride_bk

        # # # 3rd iteration
        # a = tl.load(
        #     a_ptrs,
        #     mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
        #     other=0.0,
        # )
        # b = tl.load(
        #     b_ptrs,
        #     mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 2) * BLOCK_SIZE_K,
        #     other=0.0,
        # )

        # k_start = k_start + BLOCK_SIZE_K
        # offs_ks = k_start // group_k
        # a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
        # b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

        # accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        # a_ptrs += BLOCK_SIZE_K * stride_ak
        # b_ptrs += BLOCK_SIZE_K * stride_bk

        # # 4th iteration
        # a = tl.load(
        #     a_ptrs,
        #     mask=offs_k[None, :] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
        #     other=0.0,
        # )
        # b = tl.load(
        #     b_ptrs,
        #     mask=offs_k[:, None] < K - (k * UNROLL_FACTOR + 3) * BLOCK_SIZE_K,
        #     other=0.0,
        # )

        # k_start = k_start + BLOCK_SIZE_K
        # offs_ks = k_start // group_k
        # a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
        # b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

        # accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        # a_ptrs += BLOCK_SIZE_K * stride_ak
        # b_ptrs += BLOCK_SIZE_K * stride_bk

    if C.dtype.element_ty == tl.bfloat16:
        c = accumulator.to(tl.bfloat16)
    elif C.dtype.element_ty == tl.float16:
        c = accumulator.to(tl.float16)
    else:
        c = accumulator.to(tl.float32)

    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

def triton_custom_kernel(data: input_t) -> output_t:
    """This function performs matrix multiplication with block-wise quantization.

    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.

    Args:
        A: The input tensor, e.g., activation.
        B: The input tensor, e.g., weight.
        As: The per-token-group quantization scale for `A`.
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128].
        output_dytpe: The dtype of the returned tensor.

    Returns:
        torch.Tensor: The result of matmul.
    """
    A, B, As, Bs, C = data
    
    # A = A.contiguous()
    # B = B.contiguous() 

    block_size = (128, 128)
    
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]
    
    M = A.numel() // A.shape[-1]

    assert B.ndim == 2 and Bs.ndim == 2
    N, K = B.shape
    assert triton.cdiv(N, block_n) == Bs.shape[0]
    assert triton.cdiv(K, block_k) == Bs.shape[1]


    configs = None
    def grid(META):
        return (
            triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
        )

    kernel = _w8a8_block_fp8_matmul_unrolledx4

    kernel[grid](
        A,
        B,
        C,
        As,
        Bs,
        M,
        N,
        K,
        block_n,
        block_k,
        A.stride(-2),
        A.stride(-1),
        B.stride(1),
        B.stride(0),
        C.stride(-2),
        C.stride(-1),
        As.stride(-2),
        As.stride(-1),
        Bs.stride(1),
        Bs.stride(0),
        # **config,
    )

    return C




CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

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

#include <hip/hip_fp8.h>

#include <iostream>
#include <iomanip>

using half = __half;
using float16_t = _Float16;

using fp8_t = __hip_fp8_e4m3_fnuz;

#include <random>

#include <rocwmma/rocwmma.hpp>

#define FRAG_M 32
#define FRAG_N 32
#define FRAG_K 16

#define NUM_REPEAT_MN (4)

#define BLOCK_SIZE_M (FRAG_M * NUM_REPEAT_MN)
#define BLOCK_SIZE_N (FRAG_N * NUM_REPEAT_MN)
#define BLOCK_SIZE_K (FRAG_K * 8)

#define USE_ASM false
#define ALU_ASYN false

__device__ __inline__ void __sync_warp(void) { __asm__ volatile("s_barrier" : : : "memory"); }


constexpr const int BLOCK = 128;

__global__ void v_mfma_intrinsics_test(const fp8_t* __restrict__ lhs, const fp8_t* __restrict__ rhs, __hip_bfloat16* output,
                                       unsigned int m, unsigned int n, unsigned int k,
                                       unsigned int lda, unsigned int ldb, unsigned int ldc,
                                       const float* as, const float* bs) {
    int lane_id = threadIdx.x % warpSize;
    constexpr unsigned int ele_per_thread = FRAG_M * FRAG_N / warpSize;

    int c_row_base = (blockIdx.x * BLOCK_SIZE_M);
    int c_col_base = (blockIdx.y * BLOCK_SIZE_N);

    using fp8_t = __hip_fp8_e4m3_fnuz;

    using uint_v2 = __attribute__((__vector_size__(2 * sizeof(uint)))) uint;

    __shared__ int2 lhs_tile[BLOCK_SIZE_K / 8][BLOCK_SIZE_M];
    __shared__ int2 rhs_tile[BLOCK_SIZE_K / 8][BLOCK_SIZE_N];

    // fp8_t* lhs_tile_fp8_base_ptr = &(lhs_tile_fp8[0][0]);
    int2* lhs_tile_base_ptr = (int2*)(&(lhs_tile[0][0]));
    int2* rhs_tile_base_ptr = &(rhs_tile[0][0]);
    // fp8_t* rhs_tile_base_ptr = &(rhs_tile[0][0]);

    // the output fragment is stored in 4 x AccVGPRs (see CDNA3 ISA), each thread process 4 fp32 elements
    using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
    uint result_frag[2][16] = {0};

    int sn = (n + BLOCK - 1) / BLOCK;
    float as_v[2][ele_per_thread];


    int thread_id_in_a_block = threadIdx.x + threadIdx.y * blockDim.x;

    for (int bi=0; bi < k; bi+= BLOCK) {
        float bs_v = bs[bi / BLOCK * sn + c_col_base / BLOCK];

        uint acc_frag[2][16] = {0};
        for(int ri = 0; ri < 2; ++ri) {
            for (int jj=0; jj < ele_per_thread ; jj++) {
                int j = ri * ele_per_thread + jj;
                auto x = ( lane_id / 32 ) % 2;
                auto z = (j / 4) * 8 + j % 4;
                auto inIdx = bi / BLOCK * m + (z + x * 4) + (threadIdx.y / 4 * 64) + c_row_base;
                as_v[ri][jj] = as[inIdx];
            }
        }

        constexpr int thr_per_row = 16; // 64 / 4 = 16
        if(threadIdx.y < 4)
        {
            constexpr int vec_size = 8; // 128 / (64 / 4) = 8
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            // m dim
            int g_col = c_row_base + s_col * vec_size;
            // k dim
            int g_row = bi + s_row * vec_size;

            fp8_t tmp_fp8[vec_size * vec_size];
            fp8_t tmp2_fp8[vec_size * vec_size];

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* lhs_2i_ptr = (int2 *)(lhs + (g_row + kk) * lda + g_col);
                auto ptr = tmp_fp8 + kk * vec_size;
                int2* ptr_2i_ptr = (int2*)ptr;
                *ptr_2i_ptr = *lhs_2i_ptr;
            }

            for(int tr = 0; tr < vec_size; ++tr) {
                for(int tc = 0; tc < vec_size; ++tc) {
                    tmp2_fp8[tc * vec_size + tr] = tmp_fp8[tr * vec_size + tc];
                }
            }

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* lhs_tile_2i_ptr = (int2 *)(lhs_tile_base_ptr + s_row * BLOCK_SIZE_M + s_col * 8 + kk);
                *lhs_tile_2i_ptr = *((int2*)(tmp2_fp8 + kk * vec_size));
            }

            // for debug.
            // if(threadIdx.x == 0 && threadIdx.y == 0) {
            //     fp8_t* ptr = (fp8_t*)(&(lhs_tile[0][0]));
            //     for(int pk = 0; pk < 128; ++pk){ 
            //         for(int pm = 0; pm < 128; ++pm) {
            //             int idx = pk * 128 + pm;
            //             ptr[idx] = 0.5;
            //         }
            //     }
            // }
        } // load_tile_a
        else
        {
            thread_id_in_a_block = thread_id_in_a_block % 256;
            constexpr int vec_size = 8;
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            int g_col = c_col_base + s_col * vec_size;
            int g_row = bi + s_row * vec_size;

            fp8_t tmp_fp8[vec_size * vec_size];
            fp8_t tmp2_fp8[vec_size * vec_size];

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* rhs_2i_ptr = (int2*)(rhs + (g_row + kk) * ldb + g_col);
                auto ptr = tmp_fp8 + kk * vec_size;
                int2* ptr_2i_ptr = (int2*)ptr;
                *ptr_2i_ptr = *rhs_2i_ptr;     
            }

            for(int tr = 0; tr < vec_size; ++tr) {
                for(int tc = 0; tc < vec_size; ++tc) {
                    tmp2_fp8[tc * vec_size + tr] = tmp_fp8[tr * vec_size + tc];
                }
            }
            for(int kk = 0; kk < vec_size; ++kk) {
                int2* rhs_tile_2i_ptr = (int2 *)(rhs_tile_base_ptr + s_row * BLOCK_SIZE_N + s_col * 8 + kk);
                *rhs_tile_2i_ptr = *((int2*)(tmp2_fp8 + kk * vec_size));
            }
        } // load_tile_b

        __syncthreads();


        for(int ii = 0; ii < BLOCK; ii += 16) {
            int i = bi + ii;

            fp8_t v_a_frag[2][8] = {0};
            fp8_t b_frag[8] = {0};

            int2* lhs_tile_ptr = nullptr;
            int2* rhs_tile_ptr = nullptr;
            {
                int s_row = (lane_id / 32);
                int s_col = (lane_id % 32);
                int warp_col = threadIdx.y / 4;
                lhs_tile_ptr = lhs_tile_base_ptr +(s_row + ii / 8) * BLOCK_SIZE_M + s_col + warp_col * 64;

                for(int ri = 0; ri < 2; ++ri) {
                    int2 * ptr = (int2 *)v_a_frag[ri];
                    *ptr = lhs_tile_ptr[ri * 32];
                }
            } //  compute smem window for compute
    
            {
                int s_row = (lane_id / 32);
                int s_col = (lane_id % 32);
                int warp_col = threadIdx.y % 4;
                rhs_tile_ptr = rhs_tile_base_ptr + (s_row + ii / 8) * BLOCK_SIZE_N + s_col + warp_col * 32;

                int2* b_ptr = (int2 *)b_frag;
                *b_ptr = * rhs_tile_ptr;
            } // compute smem window for compute
    
            for(int ri = 0; ri < 2; ++ri) {
                // fp8_t a_frag[8] = {0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5};
                // fp8_t b_frag[8] = {0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5};
                auto a_frag = v_a_frag[ri];

                long *a = reinterpret_cast<long *>(a_frag);
                long *b = reinterpret_cast<long *>(b_frag);
        
                // floatx16 *d = reinterpret_cast<floatx16 *>(result_frag[ri * 2 + rj]);
                floatx16 *d = reinterpret_cast<floatx16 *>(acc_frag[ri]);
                __sync_warp();
        
                *d = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                    *a,
                    *b,
                    *d, 0, 0, 0);
            }
        }

        {
            #pragma unroll
            for(int ri = 0; ri < 2; ++ri) {
                floatx16 *a = reinterpret_cast<floatx16 *>(as_v[ri]);
                floatx16 *d = reinterpret_cast<floatx16 *>(acc_frag[ri]);
                floatx16 *r = reinterpret_cast<floatx16 *>(result_frag[ri]);

                auto tmp = bs_v * (*d) * (*a);
                (*r) += tmp;
            }
        }
    }

    // Save The Result.
    __hip_bfloat16* base_output_tile_ptr = output + c_row_base * ldc + c_col_base + (threadIdx.y / 4) * 64 * ldc + (threadIdx.y % 4) * 32;

    for(int ri = 0; ri < 2; ++ri) {
        __hip_bfloat16* tmp_output_ptr = base_output_tile_ptr + ri * 32 * ldc;
        for(int rj = 0; rj < 1; ++rj) {
            tmp_output_ptr = tmp_output_ptr + rj * 32;
            floatx16 *d = reinterpret_cast<floatx16 *>(result_frag[ri]);
            for (int j=0; j < ele_per_thread; j++) {
                auto x = ( lane_id / 32 ) % 2;
                auto y = lane_id % 32; // y + j * 16 is the output data lane ID
                auto z = (j / 4) * 8 + j % 4;
                auto outIdx = y + (z + x * 4) * ldc;
                tmp_output_ptr[outIdx] = (__hip_bfloat16)((*d)[j]);

                // tmp_output_ptr[outIdx] = thread_id_in_a_block;
                // tmp_output_ptr[outIdx] = threadIdx.y;
            }
        }
    }
}




__global__ void v_mfma_intrinsics_test_splitk(const fp8_t* __restrict__ lhs, const fp8_t* __restrict__ rhs, float* output,
                                       unsigned int m, unsigned int n, unsigned int k,
                                       unsigned int lda, unsigned int ldb, unsigned int ldc,
                                       const float* as, const float* bs, int splitk) {
    int lane_id = threadIdx.x % warpSize;
    constexpr unsigned int ele_per_thread = FRAG_M * FRAG_N / warpSize;

    int k_size = k / splitk;
    int start_k = blockIdx.z * k_size;
    int end_k = start_k + k_size;

    int c_row_base = (blockIdx.x * BLOCK_SIZE_M);
    int c_col_base = (blockIdx.y * BLOCK_SIZE_N);

    using fp8_t = __hip_fp8_e4m3_fnuz;
    using uint_v2 = __attribute__((__vector_size__(2 * sizeof(uint)))) uint;

    __shared__ int2 lhs_tile[BLOCK_SIZE_K / 8][BLOCK_SIZE_M];
    __shared__ int2 rhs_tile[BLOCK_SIZE_K / 8][BLOCK_SIZE_N];

    int2* lhs_tile_base_ptr = &(lhs_tile[0][0]);
    int2* rhs_tile_base_ptr = &(rhs_tile[0][0]);

    // the output fragment is stored in 4 x AccVGPRs (see CDNA3 ISA), each thread process 4 fp32 elements
    using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
    uint result_frag[2][16] = {0};

    int sn = (n + BLOCK - 1) / BLOCK;
    float as_v[2][ele_per_thread];


    int thread_id_in_a_block = threadIdx.x + threadIdx.y * blockDim.x;

    for (int bi = start_k; bi < end_k; bi+= BLOCK) {
        float bs_v = bs[bi / BLOCK * sn + c_col_base / BLOCK];

        uint acc_frag[2][16] = {0};
        for(int ri = 0; ri < 2; ++ri) {
            for (int jj=0; jj < ele_per_thread ; jj++) {
                int j = ri * ele_per_thread + jj;
                auto x = ( lane_id / 32 ) % 2;
                auto z = (j / 4) * 8 + j % 4;
                auto inIdx = bi / BLOCK * m + (z + x * 4) + (threadIdx.y / 4 * 64) + c_row_base;
                as_v[ri][jj] = as[inIdx];
            }
        }

        constexpr int thr_per_row = 16; // 64 / 4 = 16
        if(threadIdx.y < 4)
        {
            constexpr int vec_size = 8; // 128 / (64 / 4) = 8
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            // m dim
            int g_col = c_row_base + s_col * vec_size;
            // k dim
            int g_row = bi + s_row * vec_size;

            fp8_t tmp_fp8[vec_size * vec_size];
            fp8_t tmp2_fp8[vec_size * vec_size];

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* lhs_2i_ptr = (int2 *)(lhs + (g_row + kk) * lda + g_col);
                auto ptr = tmp_fp8 + kk * vec_size;
                int2* ptr_2i_ptr = (int2*)ptr;
                *ptr_2i_ptr = *lhs_2i_ptr;
            }

            for(int tr = 0; tr < vec_size; ++tr) {
                for(int tc = 0; tc < vec_size; ++tc) {
                    tmp2_fp8[tc * vec_size + tr] = tmp_fp8[tr * vec_size + tc];
                }
            }

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* lhs_tile_2i_ptr = (int2 *)(lhs_tile_base_ptr + s_row * BLOCK_SIZE_M + s_col * 8 + kk);
                *lhs_tile_2i_ptr = *((int2*)(tmp2_fp8 + kk * vec_size));
            }

            // for debug.
            // if(threadIdx.x == 0 && threadIdx.y == 0) {
            //     fp8_t* ptr = (fp8_t*)(&(lhs_tile[0][0]));
            //     for(int pk = 0; pk < 128; ++pk){ 
            //         for(int pm = 0; pm < 128; ++pm) {
            //             int idx = pk * 128 + pm;
            //             ptr[idx] = 0.5;
            //         }
            //     }
            // }
        } // load_tile_a
        else
        {
            thread_id_in_a_block = thread_id_in_a_block % 256;
            constexpr int vec_size = 8;
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            int g_col = c_col_base + s_col * vec_size;
            int g_row = bi + s_row * vec_size;

            fp8_t tmp_fp8[vec_size * vec_size];
            fp8_t tmp2_fp8[vec_size * vec_size];

            for(int kk = 0; kk < vec_size; ++kk) {
                int2* rhs_2i_ptr = (int2*)(rhs + (g_row + kk) * ldb + g_col);
                auto ptr = tmp_fp8 + kk * vec_size;
                int2* ptr_2i_ptr = (int2*)ptr;
                *ptr_2i_ptr = *rhs_2i_ptr;     
            }

            for(int tr = 0; tr < vec_size; ++tr) {
                for(int tc = 0; tc < vec_size; ++tc) {
                    tmp2_fp8[tc * vec_size + tr] = tmp_fp8[tr * vec_size + tc];
                }
            }
            for(int kk = 0; kk < vec_size; ++kk) {
                int2* rhs_tile_2i_ptr = (int2 *)(rhs_tile_base_ptr + s_row * BLOCK_SIZE_N + s_col * 8 + kk);
                *rhs_tile_2i_ptr = *((int2*)(tmp2_fp8 + kk * vec_size));
            }
        } // load_tile_b

        __syncthreads();


        for(int ii = 0; ii < BLOCK; ii += 16) {
            int i = bi + ii;

            fp8_t v_a_frag[2][8] = {0};
            fp8_t b_frag[8] = {0};

            int2* lhs_tile_ptr = nullptr;
            int2* rhs_tile_ptr = nullptr;
            {
                int s_row = (lane_id / 32);
                int s_col = (lane_id % 32);
                int warp_col = threadIdx.y / 4;
                lhs_tile_ptr = lhs_tile_base_ptr +(s_row + ii / 8) * BLOCK_SIZE_M + s_col + warp_col * 64;

                for(int ri = 0; ri < 2; ++ri) {
                    int2 * ptr = (int2 *)v_a_frag[ri];
                    *ptr = lhs_tile_ptr[ri * 32];
                }
            } //  compute smem window for compute
    
            {
                int s_row = (lane_id / 32);
                int s_col = (lane_id % 32);
                int warp_col = threadIdx.y % 4;
                rhs_tile_ptr = rhs_tile_base_ptr + (s_row + ii / 8) * BLOCK_SIZE_N + s_col + warp_col * 32;

                int2* b_ptr = (int2 *)b_frag;
                *b_ptr = * rhs_tile_ptr;
            } // compute smem window for compute
    
            for(int ri = 0; ri < 2; ++ri) {
                // fp8_t a_frag[8] = {0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5};
                // fp8_t b_frag[8] = {0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5};
                auto a_frag = v_a_frag[ri];

                long *a = reinterpret_cast<long *>(a_frag);
                long *b = reinterpret_cast<long *>(b_frag);
        
                // floatx16 *d = reinterpret_cast<floatx16 *>(result_frag[ri * 2 + rj]);
                floatx16 *d = reinterpret_cast<floatx16 *>(acc_frag[ri]);
                __sync_warp();
        
                *d = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                    *a,
                    *b,
                    *d, 0, 0, 0);
            }
        }

        {
            #pragma unroll
            for(int ri = 0; ri < 2; ++ri) {
                floatx16 *a = reinterpret_cast<floatx16 *>(as_v[ri]);
                floatx16 *d = reinterpret_cast<floatx16 *>(acc_frag[ri]);
                floatx16 *r = reinterpret_cast<floatx16 *>(result_frag[ri]);

                auto tmp = bs_v * (*d) * (*a);
                (*r) += tmp;
            }
        }
    }

    // Save The Result.
    float* base_output_tile_ptr = output + c_row_base * ldc + c_col_base + (threadIdx.y / 4) * 64 * ldc + (threadIdx.y % 4) * 32;

    for(int ri = 0; ri < 2; ++ri) {
        float* tmp_output_ptr = base_output_tile_ptr + ri * 32 * ldc;
        for(int rj = 0; rj < 1; ++rj) {
            tmp_output_ptr = tmp_output_ptr + rj * 32;
            floatx16 *d = reinterpret_cast<floatx16 *>(result_frag[ri]);
            for (int j=0; j < ele_per_thread; j++) {
                auto x = ( lane_id / 32 ) % 2;
                auto y = lane_id % 32; // y + j * 16 is the output data lane ID
                auto z = (j / 4) * 8 + j % 4;
                auto outIdx = y + (z + x * 4) * ldc;
                // tmp_output_ptr[outIdx] = (*d)[j];

                atomicAdd( &(tmp_output_ptr[outIdx]), (*d)[j]);

                // tmp_output_ptr[outIdx] = thread_id_in_a_block;
                // tmp_output_ptr[outIdx] = threadIdx.y;
            }
        }
    }
}

__global__ void convert_float32_to_bf16(float *input_fp32, __hip_bfloat16 *output_bf16,  int num) {

    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < num){
        output_bf16[tid] = (__hip_bfloat16)input_fp32[tid];
    }
}


void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    int m = a.size(0);
    int n = b.size(0);
    int k = a.size(1);

    #define CEILDIV(x, y) (((x) + (y)-1) / (y))
    
    
    if((m == 1024 && n == 1536 && k == 7168) || 
        (m == 1024 && n == 512 && k == 7168)) {
        
        const int splitk = 8;
        dim3 grid(CEILDIV(m, BLOCK_SIZE_M), CEILDIV(n, BLOCK_SIZE_N), splitk);
        dim3 block(64, 8);
        
        float* output_d_float = nullptr;
        hipMalloc(&output_d_float, sizeof(float) * m * n);
        hipMemset(output_d_float, 0, sizeof(float) * m * n);
        
        v_mfma_intrinsics_test_splitk<<<grid, block>>>((__hip_fp8_e4m3_fnuz*)a.data_ptr(),
                                                (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
                                                output_d_float,
                                                m, n, k,
                                                m, n, n,
                                                as.data_ptr<float>(),
                                                bs.data_ptr<float>(),
                                                splitk);

        dim3 convert_block(512);
        dim3 convert_grid(CEILDIV(m * n, 512));

        convert_float32_to_bf16<<<convert_grid, convert_block>>>(output_d_float, (__hip_bfloat16*)c.data_ptr(), m * n);
    }
    else {
        dim3 grid(CEILDIV(m, BLOCK_SIZE_M), CEILDIV(n, BLOCK_SIZE_N));
        dim3 block(64, 8);
        
        v_mfma_intrinsics_test<<<grid, block>>>((__hip_fp8_e4m3_fnuz*)a.data_ptr(),
                                                (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
                                                (__hip_bfloat16*)c.data_ptr(),
                                                m, n, k,
                                                m, n, n,
                                                as.data_ptr<float>(),
                                                bs.data_ptr<float>());
    }


}

"""

import os
os.environ["CXX"] = "clang++"


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



def hip_custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c
    

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    
    M, K = a.shape
    N, K = b.shape
    
    if M % 128 == 0 and N % 128 == 0 and K % 128 == 0:
        return hip_custom_kernel(data)
    else:
        return triton_custom_kernel(data)
    

python eval.py benchmark /tmp/tmpd8pswi6e
[92mSuccessfully preprocessed all matching files.[0m

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


Total number of replaced kernel launches: 3
[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 fp8_mm.so
