#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpus MI300

import os
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline

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

cuda_src = r"""
#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>

__host__ __device__ __forceinline__ size_t cdiv(size_t a, size_t b) {
  return (a + b - 1) / b;
}

#define BLOCK_DIM 128

template <const size_t BN, const size_t BK, const size_t BM>
__global__ void fp8_mm_kernel(const __hip_fp8_e4m3_fnuz *A,
                              const __hip_fp8_e4m3_fnuz *B,
                              const float *A_scale, const float *B_scale,
                              __hip_bfloat16 *C, size_t N, size_t K, size_t M) {
  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;

  static constexpr size_t VECTOR_SIZE = 4;
  static constexpr size_t WARPSIZE = 64;
  static constexpr size_t numThreads = (BN * BM) / 16;
  static constexpr size_t strideA = (numThreads / (BN / VECTOR_SIZE));
  static constexpr size_t strideB = (numThreads / (BM / VECTOR_SIZE));

  size_t rowOffsetC = blockIdx.y * BN;
  size_t colOffsetC = blockIdx.x * BM;
  size_t colOffsetA = rowOffsetC;
  size_t colOffsetB = colOffsetC;

  size_t innerColA = threadIdx.x % (BN / VECTOR_SIZE);
  size_t innerRowA = threadIdx.x / (BN / VECTOR_SIZE);
  size_t innerColB = threadIdx.x % (BM / VECTOR_SIZE);
  size_t innerRowB = threadIdx.x / (BM / VECTOR_SIZE);

  size_t laneIdx = threadIdx.x % WARPSIZE;
  size_t warpIdx = threadIdx.x / WARPSIZE;
  size_t warpColOffset = (warpIdx % (BM / 32)) * 32;
  size_t warpRowOffset = (warpIdx / (BM / 32)) * 32;
  size_t warpX = laneIdx % 32;
  size_t warpY = laneIdx / 32;

  __shared__ __hip_fp8_e4m3_fnuz As[BK][BN], Bs[BK][BM];
  __shared__ float Ws[BN];

  __hip_fp8_e4m3_fnuz a[8], b[8];
  floatx16 d = {0};

  for (size_t tileOffset = 0; tileOffset < K; tileOffset += BK) {

    for (size_t innerRowOffsetA = 0; innerRowOffsetA < BK; innerRowOffsetA += strideA) {
      if ((tileOffset + innerRowOffsetA + innerRowA) < K &&
          (colOffsetA + innerColA * VECTOR_SIZE) < N) {
        __hip_fp8x4_e4m3_fnuz tmp = *(__hip_fp8x4_e4m3_fnuz *)&A[(tileOffset + innerRowOffsetA + innerRowA) * N + (colOffsetA + innerColA * VECTOR_SIZE)];
        *(__hip_fp8x4_e4m3_fnuz *)&As[innerRowOffsetA + innerRowA][innerColA * VECTOR_SIZE] = tmp;

        if (threadIdx.x < (BN / VECTOR_SIZE)) {
          float4 tmp = *(float4 *)&A_scale[((tileOffset + innerRowOffsetA + innerRowA) / BLOCK_DIM) * N + (colOffsetA + innerColA * VECTOR_SIZE)];
          *(float4 *)&Ws[threadIdx.x * VECTOR_SIZE] = tmp;
        }
      } else {
        *(__hip_fp8x4_e4m3_fnuz *)&As[innerRowOffsetA + innerRowA][innerColA * VECTOR_SIZE] = float4{0.0f, 0.0f, 0.0f, 0.0f};
      }
    }

    for (size_t innerRowOffsetB = 0; innerRowOffsetB < BK; innerRowOffsetB += strideB) {
      if ((tileOffset + innerRowOffsetB + innerRowB) < K &&
          (colOffsetB + innerColB) < M) {
        __hip_fp8x4_e4m3_fnuz tmp = *(__hip_fp8x4_e4m3_fnuz *)&B[(tileOffset + innerRowOffsetB + innerRowB) * M + (colOffsetB + innerColB * VECTOR_SIZE)];
        *(__hip_fp8x4_e4m3_fnuz *)&Bs[innerRowOffsetB + innerRowB][innerColB * VECTOR_SIZE] = tmp;

        if (threadIdx.x == BN) {
          Ws[BN] = B_scale[((tileOffset + innerRowOffsetB + innerRowB) / BLOCK_DIM) * cdiv(M, BLOCK_DIM) + ((colOffsetB + innerColB) / BLOCK_DIM)];
        }
      } else {
        *(__hip_fp8x4_e4m3_fnuz *)&Bs[innerRowOffsetB + innerRowB][innerColB * VECTOR_SIZE] = float4{0.0f, 0.0f, 0.0f, 0.0f};
      }
    }

    __syncthreads();

    for (size_t BKOffset = 0; BKOffset < BK; BKOffset += 16) {
      for (size_t i = 0; i < 8; ++i) {
        a[i] = As[BKOffset + warpY * 8 + i][warpRowOffset + warpX];
        b[i] = Bs[BKOffset + warpY * 8 + i][warpColOffset + warpX];
      }

      floatx16 c = {0};
      c = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
          *reinterpret_cast<long *>(a), *reinterpret_cast<long *>(b), c, 0, 0, 0);

      float b_scale = Ws[BN];
      for (size_t j = 0; j < 4; ++j) {
        for (size_t i = 0; i < 4; ++i) {
          d[i + j * 4] += c[i + j * 4] * Ws[warpRowOffset + j * 8 + warpY * 4 + i] * b_scale;
        }
      }
    }

    __syncthreads();
  }

  for (size_t j = 0; j < 4; ++j) {
    for (size_t i = 0; i < 4; ++i) {
      if ((rowOffsetC + warpRowOffset + j * 8 + warpY * 4 + i) < N &&
          (colOffsetC + warpColOffset + warpX) < M) {
        C[(rowOffsetC + warpRowOffset + j * 8 + warpY * 4 + i) * M + (colOffsetC + warpColOffset + warpX)] = (__hip_bfloat16)d[i + j * 4];
      }
    }
  }
}

at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C) {
  size_t N = A.size(0), K = A.size(1), M = B.size(0);
  const size_t BK = 16, BN = 64, BM = 64;

  dim3 numThreads((BN * BM) / 16);
  dim3 numBlocks(cdiv(M, BM), cdiv(N, BN));
  fp8_mm_kernel<BN, BK, BM><<<numBlocks, numThreads>>>(
      (__hip_fp8_e4m3_fnuz *)A.data_ptr(), (__hip_fp8_e4m3_fnuz *)B.data_ptr(),
      A_scale.data_ptr<float>(), B_scale.data_ptr<float>(),
      (__hip_bfloat16 *)C.data_ptr(), N, K, M);
  return C;
}
"""

cpp_src = r"""
at::Tensor fp8_mm(at::Tensor A, at::Tensor B, at::Tensor A_scale,
                  at::Tensor B_scale, at::Tensor C);
"""

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

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

# Geometric mean latency reference: ~362 microseconds

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: 1
[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 -O3 --offload-arch=gfx942 -std=c++20 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:77:11: warning: array index 64 is past the end of the array (that has type 'float[64]') [-Warray-bounds]
   77 |           Ws[BN] = B_scale[((tileOffset + innerRowOffsetB + innerRowB) / BLOCK_DIM) * cdiv(M, BLOCK_DIM) + ((colOffsetB + innerColB) / BLOCK_DIM)];
      |           ^  ~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:124:23: note: in instantiation of function template specialization 'fp8_mm_kernel<64UL, 16UL, 64UL>' requested here
  124 |  hipLaunchKernelGGL(( fp8_mm_kernel<BN, BK, BM>), dim3(numBlocks), dim3(numThreads), 0, 0, 
      |                       ^
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:48:3: note: array 'Ws' declared here
   48 |   __shared__ float Ws[BN];
      |   ^
/opt/rocm-6.3.1/lib/llvm/lib/clang/18/include/__clang_hip_runtime_wrapper.h:24:20: note: expanded from macro '__shared__'
   24 | #define __shared__ __attribute__((shared))
      |                    ^
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:96:23: warning: array index 64 is past the end of the array (that has type 'float[64]') [-Warray-bounds]
   96 |       float b_scale = Ws[BN];
      |                       ^  ~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:48:3: note: array 'Ws' declared here
   48 |   __shared__ float Ws[BN];
      |   ^
/opt/rocm-6.3.1/lib/llvm/lib/clang/18/include/__clang_hip_runtime_wrapper.h:24:20: note: expanded from macro '__shared__'
   24 | #define __shared__ __attribute__((shared))
      |                    ^
2 warnings generated when compiling for gfx942.
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:77:11: warning: array index 64 is past the end of the array (that has type 'float[64]') [-Warray-bounds]
   77 |           Ws[BN] = B_scale[((tileOffset + innerRowOffsetB + innerRowB) / BLOCK_DIM) * cdiv(M, BLOCK_DIM) + ((colOffsetB + innerColB) / BLOCK_DIM)];
      |           ^  ~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:124:23: note: in instantiation of function template specialization 'fp8_mm_kernel<64UL, 16UL, 64UL>' requested here
  124 |  hipLaunchKernelGGL(( fp8_mm_kernel<BN, BK, BM>), dim3(numBlocks), dim3(numThreads), 0, 0, 
      |                       ^
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:48:3: note: array 'Ws' declared here
   48 |   __shared__ float Ws[BN];
      |   ^
/opt/rocm-6.3.1/lib/llvm/lib/clang/18/include/__clang_hip_runtime_wrapper.h:24:20: note: expanded from macro '__shared__'
   24 | #define __shared__ __attribute__((shared))
      |                    ^
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:96:23: warning: array index 64 is past the end of the array (that has type 'float[64]') [-Warray-bounds]
   96 |       float b_scale = Ws[BN];
      |                       ^  ~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:48:3: note: array 'Ws' declared here
   48 |   __shared__ float Ws[BN];
      |   ^
/opt/rocm-6.3.1/lib/llvm/lib/clang/18/include/__clang_hip_runtime_wrapper.h:24:20: note: expanded from macro '__shared__'
   24 | #define __shared__ __attribute__((shared))
      |                    ^
2 warnings generated when compiling for host.
[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 fp8_mm.so

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