# This script provides a template for using load_inline to run a HIP kernel for
import torch
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, torch::Tensor aT, torch::Tensor bT);
"""

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


constexpr const int BLOCKM = 64;
constexpr const int BLOCKN = 64;
constexpr const int BLOCK = 128;

using fp8_raw_t = uint8_t;
using fp8x8_t  = fp8_raw_t __attribute((ext_vector_type(8)));

using fp32x2_t = float __attribute((ext_vector_type(2)));
using fp32x4_t = float __attribute((ext_vector_type(4)));
using fp32x16_t = float __attribute((ext_vector_type(16)));

template <typename Y, typename X>
__device__ constexpr Y bit_cast(const X& x)
{
    return __builtin_bit_cast(Y, x);
}

template <typename __hip_fp8_e4m3_fnuz>
__global__ void permute_kernel(__hip_fp8_e4m3_fnuz *output, 
                               const __hip_fp8_e4m3_fnuz *input, 
                               const int mn,
                               const int k) {
    __shared__ __hip_fp8_e4m3_fnuz shbuf[32 * (32 + 1)]; 
    const int32_t warp_id  = threadIdx.y;     
    const int32_t lane_id  = threadIdx.x;     
    const int32_t k_id  = blockIdx.y * 32;
    const int32_t mn_id = blockIdx.x * 32;

    const size_t input_idx = (k_id + warp_id) * mn + mn_id + lane_id;

    const int lane_id_x_33 = lane_id * 33;
    #pragma unroll
    for (int kLoopIdx = 0; kLoopIdx < 4; kLoopIdx++) {
      shbuf[lane_id_x_33 + warp_id + kLoopIdx * 8] = input[input_idx + kLoopIdx * 8 * mn];
    }
    __syncthreads();

    const int32_t kiOut = k_id + lane_id;
    int mnI = warp_id;
    #pragma unroll
    for (int mnLoopIdx = 0; mnLoopIdx < 4; ++mnLoopIdx) {
      output[kiOut + (mn_id + mnI) * k] = shbuf[(mnI)*33 + lane_id];
      mnI += 8;
    }
}

__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz *a, const __hip_fp8_e4m3_fnuz *b, const float *as, const float *bs,
                              __hip_bfloat16 *c, int m, int n, int k)
{

  // Your implementation here
  int bx = blockIdx.x;
  int by = blockIdx.y;

  int tid = threadIdx.x;
  int lane_id = threadIdx.x % 64;
  int warp_id = threadIdx.x / 64;
  int lane_mod32 = lane_id % 32;
  int lane_div32 = lane_id / 32;
  int lane_mod16 = lane_id % 16;
  int lane_div16 = lane_id / 16;
  int lane_mod8 = lane_id % 8;
  int lane_div8 = lane_id / 8;
  int warp_mod2 = warp_id % 2;
  int warp_div2 = warp_id / 2;

  int gemm_rowA = bx * BLOCKM + warp_id * 16 + lane_div8;
  int gemm_col = lane_mod8 * 16;

  int gemm_rowB = by * BLOCKN + warp_id * 16 + lane_div8;

  int A_ldg_idx = gemm_rowA * k + gemm_col;
  int B_ldg_idx = gemm_rowB * k + gemm_col;

  int A_sts_idx = (warp_id * 16 + lane_div8) * 128 + (lane_mod8 ^ lane_div8) * 16;
  int B_sts_idx = (warp_id * 16 + lane_div8) * 128 + (lane_mod8 ^ lane_div8) * 16;
  
  int lds_row = lane_mod32;

  int A_lds_idx[4];
  int B_lds_idx[4];

  for (int i = 0; i < 4; i++){
    int lds_col = ((lane_div32 + 2 * i) % 8) ^ (lds_row % 8);
    A_lds_idx[i] = (warp_div2 * 32 * 128) + lds_row * 128 + lds_col * 16;
    B_lds_idx[i] = (warp_mod2 * 32 * 128) + lds_row * 128 + lds_col * 16;
  }

  int sn = (n + BLOCK - 1) / BLOCK;

  as += BLOCKM * blockIdx.x;
  bs += BLOCKN * blockIdx.y / BLOCK;
  c += BLOCKN * blockIdx.y + BLOCKM * blockIdx.x * n;
  // smem
  __shared__ char smem[((BLOCKM + BLOCKN) * BLOCK) * sizeof(__hip_fp8_e4m3_fnuz) + (1 + BLOCKM) * sizeof(float)];
  __hip_fp8_e4m3_fnuz *smemA = reinterpret_cast<__hip_fp8_e4m3_fnuz *>(smem);
  __hip_fp8_e4m3_fnuz *smemB = reinterpret_cast<__hip_fp8_e4m3_fnuz *>(smem + BLOCKM * BLOCK * sizeof(__hip_fp8_e4m3_fnuz));
  float *smem_scales_a = reinterpret_cast<float *>(smem + (BLOCKM + BLOCKN) * BLOCK * sizeof(__hip_fp8_e4m3_fnuz));
  float *smem_scales_b = reinterpret_cast<float *>(smem + ((BLOCKM + BLOCKN) * BLOCK) * sizeof(__hip_fp8_e4m3_fnuz) + BLOCKM * sizeof(float));

  fp32x16_t acc[1][1];
  for (int i = 0; i < 1; i++)
  {
    for (int j = 0; j < 1; j++)
    {
      acc[i][j] = {0.0f};
    }
  }

  fp32x4_t A_ldg_reg[2];
  fp32x4_t B_ldg_reg[2];

  __hip_fp8_e4m3_fnuz A_lds_reg[1][2][8];
  __hip_fp8_e4m3_fnuz B_lds_reg[1][2][8];
  float A_scale_lds_reg;
  float B_scale_lds_reg;

  for (int itile = 0; itile < k; itile += BLOCK)
  {

    fp32x16_t block_acc[1][1];
    for (int i = 0; i < 1; i++)
    {
      for (int j = 0; j < 1; j++)
      {
        block_acc[i][j] = {0.0f};
      }
    }
    // ldg
    for (int i = 0; i < 2; i++)
    {
      A_ldg_reg[i] = *(fp32x4_t *)(&a[A_ldg_idx + i * 8 * k]);
      B_ldg_reg[i] = *(fp32x4_t *)(&b[B_ldg_idx + i * 8 * k]);
    }
    A_ldg_idx += BLOCK;
    B_ldg_idx += BLOCK;
    // sts
    for (int i = 0; i < 2; i++)
    {
      *(fp32x4_t *)(&smemA[A_sts_idx + i * 8 * 128]) = A_ldg_reg[i];
      *(fp32x4_t *)(&smemB[B_sts_idx + i * 8 * 128]) = B_ldg_reg[i];
    }
    if (tid < BLOCKM)
    {
      smem_scales_a[tid] = as[tid];
    }
    if (tid == 0)
    {
      smem_scales_b[0] = bs[0];
    }
    __syncthreads();


    A_scale_lds_reg = smem_scales_a[lane_mod32 + warp_div2 * 32];
    B_scale_lds_reg = smem_scales_b[0];

    for (int k_step = 0; k_step < 4; k_step++)
    {
      // lds
      for (int i = 0; i < 1; i++)
      {
        *(fp32x4_t *)(&A_lds_reg[i]) = *(fp32x4_t *)(&smemA[A_lds_idx[k_step] + i * 16 * 128]);
      }
      for (int j = 0; j < 1; j++)
      {
        *(fp32x4_t *)(&B_lds_reg[j]) = *(fp32x4_t *)(&smemB[B_lds_idx[k_step] + j * 16 * 128]);
      }

      for (int i = 0; i < 1; i++)
      {
        for (int j = 0; j < 1; j++)
        {
          block_acc[i][j] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
              bit_cast<long>(B_lds_reg[i][0]), bit_cast<long>(A_lds_reg[j][0]), block_acc[i][j], 0, 0, 0);
          block_acc[i][j] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
              bit_cast<long>(B_lds_reg[i][1]), bit_cast<long>(A_lds_reg[j][1]), block_acc[i][j], 0, 0, 0);
        }
      }
    }

    for (int i = 0; i < 1; ++i)
    {
      for (int j = 0; j < 1; ++j)
      {
        for (int k = 0; k < 16; ++k)
        {
          acc[i][j][k] += block_acc[i][j][k] * A_scale_lds_reg * B_scale_lds_reg;
        }
      }
    }
    as += m;
    bs += sn;
    __syncthreads();
  }

  // finally, write the result as bf16
  for (int i = 0; i < 1; ++i)
  {
    for (int j = 0; j < 1; ++j)
    {
      for (int k = 0; k < 4; ++k)
      {
        int C_stg_idx = (lane_mod32 + warp_div2 * 32 + 16 * j) * n + (lane_div32 * 4 + k * 8 + warp_mod2 * 32 + 16 * i);
        __hip_bfloat16 temp[4];
        temp[0] = (__hip_bfloat16)acc[i][j][k * 4 + 0];
        temp[1] = (__hip_bfloat16)acc[i][j][k * 4 + 1];
        temp[2] = (__hip_bfloat16)acc[i][j][k * 4 + 2];
        temp[3] = (__hip_bfloat16)acc[i][j][k * 4 + 3];

        *(fp32x2_t *)(&c[C_stg_idx]) = *(fp32x2_t *)(&temp[0]);
      }
    }
  }
}


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

    permute_kernel<<<dim3((m+32-1)/32, (k+32-1)/32), dim3(32, 8), 0, 0>>> ((__hip_fp8_e4m3_fnuz*)aT.data_ptr(), (__hip_fp8_e4m3_fnuz*)a.data_ptr(), m, k);
    permute_kernel<<<dim3((n+32-1)/32, (k+32-1)/32), dim3(32, 8), 0, 0>>> ((__hip_fp8_e4m3_fnuz*)bT.data_ptr(), (__hip_fp8_e4m3_fnuz*)b.data_ptr(), n, k);
    custom_kernel<<<dim3((m+BLOCKM-1)/BLOCKM, (n+BLOCKN-1)/BLOCKN), dim3(256), 0, 0>>> ((__hip_fp8_e4m3_fnuz*)aT.data_ptr(), (__hip_fp8_e4m3_fnuz*)bT.data_ptr(), 
    // ori_kernel<<<dim3((m+15)/16, (n+15)/16), dim3(16, 16), 0, 0>>> ((__hip_fp8_e4m3_fnuz*)aT.data_ptr(), (__hip_fp8_e4m3_fnuz*)bT.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", "-save-temps"],
)


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    aT = (torch.empty((k, m), dtype=torch.bfloat16, device="cuda")).to(torch.float8_e4m3fnuz)
    bT = (torch.empty((k, n), dtype=torch.bfloat16, device="cuda")).to(torch.float8_e4m3fnuz)

    module.fp8_mm(a, b, a_scale, b_scale, c, aT, bT)
    return c


python submission.py
[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 [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 3
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=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/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=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 -std=c++20 -save-temps -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] 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 fp8_mm.so

python eval.py test /tmp/tmpc6eoe78a
[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
ninja: no work to do.
