# This script provides a template for using load_inline to run a HIP kernel for
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
CPP_WRAPPER = """
void optimized_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>
#include <stdio.h>

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

constexpr const int BLOCK = 128;

#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 * 1)

#define CEILDIV(x, y) (((x) + (y)-1) / (y))

using fp8_t = __hip_fp8_e4m3_fnuz;
using hip_bf16 = __hip_bfloat16;


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

__global__ void fp8_gemm_core(
    const fp8_t* __restrict__ lhs, const fp8_t* __restrict__ rhs, hip_bf16* 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 hip_bf16 = __hip_bfloat16;

    __shared__ fp8_t lhs_tile[BLOCK_SIZE_K][BLOCK_SIZE_M];
    __shared__ fp8_t rhs_tile[BLOCK_SIZE_K][BLOCK_SIZE_N];

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

    using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
    uint result_frag[4][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[4][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 / 2 * 64) + c_row_base;
                as_v[ri][jj] = as[inIdx];
            }
        }



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

            constexpr int thr_per_row = 16; // 64 / 4 = 16
            {
    
                constexpr int vec_size = 8; // 128 / (64 / 4) = 8
                int s_col = (thread_id_in_a_block % thr_per_row) * vec_size;
                int s_row = (thread_id_in_a_block / thr_per_row);
    
                // m dim
                int g_col = c_row_base + s_col;
                // k dim
                int g_row = i + s_row;
    
                if (g_row < k && g_col < m) {
    
                    int2* lhs_tile_2i_ptr = (int2 *)(lhs_tile_base_ptr + s_row * BLOCK_SIZE_M + s_col);
                    int2* lhs_2i_ptr = (int2 *)(lhs + g_row * lda + g_col);
    
                    *(lhs_tile_2i_ptr) = *(lhs_2i_ptr);
                }
            }
            {
                // constexpr int vec_size = 16;
                constexpr int vec_size = 8;
                int s_col = (thread_id_in_a_block % thr_per_row) * vec_size;
                int s_row = (thread_id_in_a_block / thr_per_row);
    
                int g_col = c_col_base + s_col;
                int g_row = i + s_row;
    
                if (g_row < k && g_col < n) {
                    int2* rhs_tile_2i_ptr = (int2 *)(rhs_tile_base_ptr + s_row * BLOCK_SIZE_N + s_col);
                    int2* rhs_2i_ptr = (int2 *)(rhs + g_row * ldb + g_col);
    
                    *(rhs_tile_2i_ptr) = *(rhs_2i_ptr);
                } 
            }

            __syncthreads();
    
    
            fp8_t a_frag[2][8] = {0};
            fp8_t b_frag[2][8] = {0};
            fp8_t* lhs_tile_ptr = nullptr;
            fp8_t* rhs_tile_ptr = nullptr;
    
            {
                constexpr int vec_size = 8;
                int s_row = (lane_id / 32) * vec_size;
                int s_col = (lane_id % 32);
                int warp_col = threadIdx.y / 2;
                lhs_tile_ptr = lhs_tile_base_ptr + s_row * BLOCK_SIZE_M + s_col + warp_col * 64;
            }
    
            {
                constexpr int vec_size = 8;
                int s_row = (lane_id / 32) * vec_size;
                int s_col = (lane_id % 32);
                int warp_row = threadIdx.y % 2;
                rhs_tile_ptr = rhs_tile_base_ptr + s_row * BLOCK_SIZE_N + s_col + warp_row * 64;
            }
    
    
            for(int ri = 0; ri < 2; ++ri) {
                fp8_t* tmp_lhs_tile_ptr = lhs_tile_ptr + ri * 32;
                for(int i = 0; i < 8; ++i) {
                    a_frag[ri][i] = tmp_lhs_tile_ptr[0];
                    tmp_lhs_tile_ptr += BLOCK_SIZE_M;
                }
            }
    
            for(int rj = 0; rj < 2; ++rj) {
                fp8_t* tmp_rhs_tile_ptr = rhs_tile_ptr + rj * 32;
                for (int i=0; i < 8; i++) {
                    b_frag[rj][i] = tmp_rhs_tile_ptr[0];
                    tmp_rhs_tile_ptr += BLOCK_SIZE_N;
                }
            }
    
            for(int ri = 0; ri < 2; ++ri) {
                for(int rj = 0; rj < 2; ++rj) {
                    long *a = reinterpret_cast<long *>(a_frag[ri]);
                    long *b = reinterpret_cast<long *>(b_frag[rj]);
            
                    floatx16 *d = reinterpret_cast<floatx16 *>(acc_frag[ri * 2 + rj]);
                    __sync_warp();
            
                    *d = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
                        *a,
                        *b,
                        *d, 0, 0, 0);
                }
                        
            }
        }

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

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

    }
    hip_bf16* base_output_tile_ptr = output + c_row_base * ldc + c_col_base + (threadIdx.y / 2) * 64 * ldc + (threadIdx.y % 2) * 64;
    for(int ri = 0; ri < 2; ++ri) {
        hip_bf16* tmp_output_ptr = base_output_tile_ptr + ri * 32 * ldc;
        for(int rj = 0; rj < 2; ++rj) {
            tmp_output_ptr = tmp_output_ptr + rj * 32;
            floatx16 *d = reinterpret_cast<floatx16 *>(result_frag[ri * 2 + rj]);
            for (int j=0; j < ele_per_thread; j++) {
                auto x = ( lane_id / 32 ) % 2;
                auto y = lane_id % 32;
                auto z = (j / 4) * 8 + j % 4;
                auto outIdx = y + (z + x * 4) * ldc;
                tmp_output_ptr[outIdx] = (hip_bf16)(*d)[j];
            }
        }
    }
}


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

    dim3 grid(CEILDIV(m, BLOCK_SIZE_M), CEILDIV(n, BLOCK_SIZE_N));
    dim3 block(64, 4);
    
    fp8_gemm_core<<<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++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

module = load_inline(
    name='optimized_fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['optimized_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.optimized_fp8_mm(a, b, a_scale, b_scale, c)
    return c


python submission.py
[92mSuccessfully preprocessed all matching files.[0m

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


Total number of replaced kernel launches: 1
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=optimized_fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/optimized_fp8_mm/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=optimized_fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc --offload-arch=gfx942 -std=c++20 -c /home/runner/.cache/torch_extensions/py310_cpu/optimized_fp8_mm/hip.hip -o hip.cuda.o 
[3/3] 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 optimized_fp8_mm.so

python eval.py test /tmp/tmpp9tvzdxh
[92mSuccessfully preprocessed all matching files.[0m

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


Total number of replaced kernel launches: 1
[1/1] clang++ main.o hip.cuda.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 optimized_fp8_mm.so
