# This script provides a template for using load_inline to run a HIP kernel for
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr const int BLOCK = 128;

__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz* A, const __hip_fp8_e4m3_fnuz* B, const float* alpha, const float* beta, __hip_bfloat16* C, int A_rows, int B_rows, int B_cols) {
    using float4 = __attribute__( (__vector_size__(4 * sizeof(float)) )) float;
    float4 C_frags[4] = {0.0f, 0.0f, 0.0f, 0.0f} ;

    int thread_num =  threadIdx.y * blockDim.x + threadIdx.x;
    int lane_id = thread_num%64;
    int wf_id = thread_num/64;

    int upper_left_x = blockDim.x * blockIdx.x ;
    int upper_left_y = 4*blockDim.y * blockIdx.y ;

    int A_row = (upper_left_y + lane_id%16);
    int B_row = upper_left_x + lane_id%16;

    int sn = (B_rows + 128 - 1)/128;


    __shared__  __hip_fp8_e4m3_fnuz A_tile[32][32];
    __shared__ __hip_fp8_e4m3_fnuz B_tile[32][32];

    for (int i = 0; i < B_cols; i += 32) {
        int A_col = i + thread_num%32; 
        int B_col = A_col;

        for (int row = 0; row < 4; row ++) {
            int A_tile_row = (4 * (thread_num/32)) + row; 
            int A_tile_col = thread_num%32; 
            int A_row = ((A_tile_row) + upper_left_y) ; 
            int B_row = ((A_tile_row) + upper_left_x) ; 

            A_tile[A_tile_row][thread_num%32] = A[A_row + A_col * A_rows]; 
            B_tile[A_tile_row][thread_num%32] = B[B_row + B_col * B_rows]; 
        
        }

        __syncthreads();

        for (int subtile = 0; subtile < 8; subtile++) {
            int a_row_offset = 16 * (wf_id / 2);
            int b_row_offset = 16 * (wf_id % 2);
            int col_offset   =  (subtile * 4);

            // recover indices in A 
            //
            // row: a_row_offset + lane_id % 16 + upper_left_x (cx)
            // col: col_offset + lane_id / 16

            // recover indices in B 
            //
            // row: b_row_offset + lane_id % 16 + upper_left_y (cy)
            // col: col_offset + lane_id / 16

            float A_val = (float) A_tile[a_row_offset + lane_id % 16][col_offset + lane_id / 16];
            float B_val = (float) B_tile[b_row_offset + lane_id % 16][col_offset + lane_id /16];

            // TODO: factor out scaling and using SMEM
            float a = alpha[a_row_offset + lane_id % 16 + upper_left_y + i/128 * A_rows];
            float b = beta[(b_row_offset + lane_id % 16 + upper_left_x)/128 + i/128 * sn];

            C_frags[wf_id] = __builtin_amdgcn_mfma_f32_16x16x4f32((float) a * A_val, (float) b * B_val, C_frags[wf_id], 0, 0, 0);
          }

          __syncthreads();
        
    }

    for (int p = 0; p < 4; ++p) {
        int wf_row_offset = 16 * (wf_id / 2);
        int wf_col_offset = 16 * (wf_id % 2);
        int C_row = upper_left_y + wf_row_offset + (4 * (lane_id / 16)) + p;
        int C_col = upper_left_x + wf_col_offset + (lane_id % 16);

        if (C_row < A_rows && C_col < B_rows && wf_id < 4 && p < 4) {
            C[C_row * B_rows + C_col] = (__hip_bfloat16) C_frags[wf_id][p];
        }
    }
}

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);
    custom_kernel<<<dim3(n/32, m/32), dim3(32, 8), 0, 0>>> ((__hip_fp8_e4m3_fnuz*)a.data_ptr(), (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
    as.data_ptr<float>(), bs.data_ptr<float>(), (__hip_bfloat16*)c.data_ptr(), m, n, k);
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

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 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


python eval.py leaderboard /tmp/tmp83m7tuvz
[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: 1
[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
